chore: import upstream snapshot with attribution

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

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py

@@ -0,0 +1,708 @@
+import functools
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
+
+from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import (
+    LTX2ConnectorConfig,
+)
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+def apply_interleaved_rotary_emb(
+    x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
+) -> torch.Tensor:
+    cos, sin = freqs
+    x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)  # [B, S, C // 2]
+    x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+    return x * cos + x_rotated * sin
+
+
+def apply_split_rotary_emb(
+    x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
+) -> torch.Tensor:
+    cos, sin = freqs
+
+    x_dtype = x.dtype
+    needs_reshape = False
+    if x.ndim != 4 and cos.ndim == 4:
+        # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head)
+        # The cos/sin batch dim may only be broadcastable, so take batch size from x
+        b = x.shape[0]
+        _, h, t, _ = cos.shape
+        x = x.reshape(b, t, h, -1).transpose(1, 2)
+        needs_reshape = True
+
+    # Split last dim (2*r) into (d=2, r)
+    last = x.shape[-1]
+    if last % 2 != 0:
+        raise ValueError(
+            f"Expected x.shape[-1] to be even for split rotary, got {last}."
+        )
+    r = last // 2
+
+    # (..., 2, r)
+    split_x = x.reshape(*x.shape[:-1], 2, r)
+    first_x = split_x[..., :1, :]  # (..., 1, r)
+    second_x = split_x[..., 1:, :]  # (..., 1, r)
+
+    cos_u = cos.unsqueeze(-2)  # broadcast to (..., 1, r) against (..., 2, r)
+    sin_u = sin.unsqueeze(-2)
+
+    out = split_x * cos_u
+    first_out = out[..., :1, :]
+    second_out = out[..., 1:, :]
+
+    first_out.addcmul_(-sin_u, second_x)
+    second_out.addcmul_(sin_u, first_x)
+
+    out = out.reshape(*out.shape[:-2], last)
+
+    if needs_reshape:
+        out = out.transpose(1, 2).reshape(b, t, -1)
+
+    out = out.to(dtype=x_dtype)
+    return out
+
+
+@functools.lru_cache(maxsize=5)
+def _ltx2_connector_rope_freq_grid_np(
+    theta: float, num_pos_dims: int, dim: int
+) -> torch.Tensor:
+    # Official LTX uses NumPy float64 for double-precision RoPE frequencies.
+    n_elem = 2 * num_pos_dims
+    pow_indices = np.power(
+        theta,
+        np.linspace(0.0, 1.0, dim // n_elem, dtype=np.float64),
+    )
+    return torch.tensor(pow_indices * math.pi / 2.0, dtype=torch.float32)
+
+
+class LTX2Attention(torch.nn.Module):
+    r"""
+    Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key
+    RoPE embeddings separately for audio-to-video (a2v) and video-to-audio (v2a) cross-attention.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        kv_heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = True,
+        cross_attention_dim: Optional[int] = None,
+        out_bias: bool = True,
+        qk_norm: str = "rms_norm_across_heads",
+        norm_eps: float = 1e-6,
+        norm_elementwise_affine: bool = True,
+        rope_type: str = "interleaved",
+        apply_gated_attention: bool = False,
+        processor=None,
+    ):
+        super().__init__()
+        if qk_norm != "rms_norm_across_heads":
+            raise NotImplementedError(
+                "Only 'rms_norm_across_heads' is supported as a valid value for `qk_norm`."
+            )
+
+        self.head_dim = dim_head
+        self.inner_dim = dim_head * heads
+        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+        self.query_dim = query_dim
+        self.cross_attention_dim = (
+            cross_attention_dim if cross_attention_dim is not None else query_dim
+        )
+        self.use_bias = bias
+        self.dropout = dropout
+        self.out_dim = query_dim
+        self.heads = heads
+        self.rope_type = rope_type
+
+        self.norm_q = torch.nn.RMSNorm(
+            dim_head * heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine
+        )
+        self.norm_k = torch.nn.RMSNorm(
+            dim_head * kv_heads,
+            eps=norm_eps,
+            elementwise_affine=norm_elementwise_affine,
+        )
+        self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = torch.nn.Linear(
+            self.cross_attention_dim, self.inner_kv_dim, bias=bias
+        )
+        self.to_v = torch.nn.Linear(
+            self.cross_attention_dim, self.inner_kv_dim, bias=bias
+        )
+        self.to_gate_logits = None
+        if apply_gated_attention:
+            self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True)
+        self.to_out = torch.nn.ModuleList([])
+        self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(torch.nn.Dropout(dropout))
+
+        # Scaled dot product attention
+        self.attn = USPAttention(
+            num_heads=heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends={
+                AttentionBackendEnum.FA,
+                AttentionBackendEnum.AITER,
+                AttentionBackendEnum.TORCH_SDPA,
+                AttentionBackendEnum.SAGE_ATTN,
+                AttentionBackendEnum.SAGE_ATTN_3,
+            },
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        gate_input = hidden_states
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(encoder_hidden_states)
+        value = self.to_v(encoder_hidden_states)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if query_rotary_emb is not None:
+            if self.rope_type == "interleaved":
+                query = apply_interleaved_rotary_emb(query, query_rotary_emb)
+                key = apply_interleaved_rotary_emb(
+                    key,
+                    key_rotary_emb if key_rotary_emb is not None else query_rotary_emb,
+                )
+            elif self.rope_type == "split":
+                query = apply_split_rotary_emb(query, query_rotary_emb)
+                key = apply_split_rotary_emb(
+                    key,
+                    key_rotary_emb if key_rotary_emb is not None else query_rotary_emb,
+                )
+
+        query = query.unflatten(2, (self.heads, -1)).transpose(1, 2)
+        key = key.unflatten(2, (self.heads, -1)).transpose(1, 2)
+        value = value.unflatten(2, (self.heads, -1)).transpose(1, 2)
+
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                attention_mask = attention_mask[:, None, None, :]
+            elif attention_mask.ndim == 3:
+                attention_mask = attention_mask[:, None, :, :]
+            attention_mask = attention_mask.to(dtype=query.dtype)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+        )
+        hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if self.to_gate_logits is not None:
+            gate_logits = self.to_gate_logits(gate_input)
+            b, t, _ = hidden_states.shape
+            hidden_states = hidden_states.view(b, t, self.heads, self.head_dim)
+            hidden_states = hidden_states * (
+                2.0 * torch.sigmoid(gate_logits).unsqueeze(-1)
+            )
+            hidden_states = hidden_states.view(b, t, self.heads * self.head_dim)
+
+        hidden_states = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+        return hidden_states
+
+
+class LTX2RotaryPosEmbed1d(nn.Module):
+    """
+    1D rotary positional embeddings (RoPE) for the LTX 2.0 text encoder connectors.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base_seq_len: int = 4096,
+        theta: float = 10000.0,
+        double_precision: bool = True,
+        rope_type: str = "interleaved",
+        num_attention_heads: int = 32,
+    ):
+        super().__init__()
+        if rope_type not in ["interleaved", "split"]:
+            raise ValueError(
+                f"{rope_type=} not supported. Choose between 'interleaved' and 'split'."
+            )
+
+        self.dim = dim
+        self.base_seq_len = base_seq_len
+        self.theta = theta
+        self.double_precision = double_precision
+        self.rope_type = rope_type
+        self.num_attention_heads = num_attention_heads
+
+    def forward(
+        self,
+        batch_size: int,
+        pos: int,
+        device: Union[str, torch.device],
+        dtype: Optional[torch.dtype] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Get 1D position ids
+        grid_1d = torch.arange(pos, dtype=torch.float32, device=device)
+        # Get fractional indices relative to self.base_seq_len
+        grid_1d = grid_1d / self.base_seq_len
+        grid = grid_1d.unsqueeze(0).repeat(batch_size, 1)  # [batch_size, seq_len]
+
+        # 2. Calculate 1D RoPE frequencies
+        num_rope_elems = 2  # 1 (because 1D) * 2 (for cos, sin) = 2
+        if self.double_precision:
+            freqs = _ltx2_connector_rope_freq_grid_np(self.theta, 1, self.dim).to(
+                device=device
+            )
+        else:
+            pow_indices = torch.pow(
+                self.theta,
+                torch.linspace(
+                    start=0.0,
+                    end=1.0,
+                    steps=self.dim // num_rope_elems,
+                    dtype=torch.float32,
+                    device=device,
+                ),
+            )
+            freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32)
+
+        # 3. Matrix-vector outer product between pos ids of shape (batch_size, seq_len) and freqs vector of shape
+        # (self.dim // 2,).
+        freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs  # [B, seq_len, self.dim // 2]
+
+        # 4. Get real, interleaved (cos, sin) frequencies, padded to self.dim
+        if self.rope_type == "interleaved":
+            cos_freqs = freqs.cos().repeat_interleave(2, dim=-1)
+            sin_freqs = freqs.sin().repeat_interleave(2, dim=-1)
+
+            if self.dim % num_rope_elems != 0:
+                cos_padding = torch.ones_like(
+                    cos_freqs[:, :, : self.dim % num_rope_elems]
+                )
+                sin_padding = torch.zeros_like(
+                    sin_freqs[:, :, : self.dim % num_rope_elems]
+                )
+                cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1)
+                sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1)
+
+        elif self.rope_type == "split":
+            expected_freqs = self.dim // 2
+            current_freqs = freqs.shape[-1]
+            pad_size = expected_freqs - current_freqs
+            cos_freq = freqs.cos()
+            sin_freq = freqs.sin()
+
+            if pad_size != 0:
+                cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+                sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
+
+                cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
+                sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
+
+            # Reshape freqs to be compatible with multi-head attention
+            b = cos_freq.shape[0]
+            t = cos_freq.shape[1]
+
+            cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1)
+            sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1)
+
+            cos_freqs = torch.swapaxes(cos_freq, 1, 2)  # (B,H,T,D//2)
+            sin_freqs = torch.swapaxes(sin_freq, 1, 2)  # (B,H,T,D//2)
+
+        if dtype is not None:
+            cos_freqs = cos_freqs.to(dtype)
+            sin_freqs = sin_freqs.to(dtype)
+        return cos_freqs, sin_freqs
+
+
+class LTX2TransformerBlock1d(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        activation_fn: str = "gelu-approximate",
+        eps: float = 1e-6,
+        rope_type: str = "interleaved",
+        apply_gated_attention: bool = False,
+    ):
+        super().__init__()
+
+        self.norm1 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.attn1 = LTX2Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            kv_heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            rope_type=rope_type,
+            apply_gated_attention=apply_gated_attention,
+        )
+
+        self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.ff = FeedForward(dim, activation_fn=activation_fn)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        norm_hidden_states = self.norm1(hidden_states)
+        attn_hidden_states = self.attn1(
+            norm_hidden_states,
+            attention_mask=attention_mask,
+            query_rotary_emb=rotary_emb,
+        )
+        hidden_states = hidden_states + attn_hidden_states
+
+        norm_hidden_states = self.norm2(hidden_states)
+        ff_hidden_states = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + ff_hidden_states
+
+        return hidden_states
+
+
+class LTX2ConnectorTransformer1d(nn.Module):
+    """
+    A 1D sequence transformer for modalities such as text.
+    In LTX 2.0, this is used to process the text encoder hidden states for each of the video and audio streams.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        num_attention_heads: int = 30,
+        attention_head_dim: int = 128,
+        num_layers: int = 2,
+        num_learnable_registers: int | None = 128,
+        rope_base_seq_len: int = 4096,
+        rope_theta: float = 10000.0,
+        rope_double_precision: bool = True,
+        eps: float = 1e-6,
+        causal_temporal_positioning: bool = False,
+        rope_type: str = "interleaved",
+        apply_gated_attention: bool = False,
+    ):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.causal_temporal_positioning = causal_temporal_positioning
+
+        self.num_learnable_registers = num_learnable_registers
+        self.learnable_registers = None
+        if num_learnable_registers is not None:
+            init_registers = (
+                torch.rand(num_learnable_registers, self.inner_dim) * 2.0 - 1.0
+            )
+            self.learnable_registers = torch.nn.Parameter(init_registers)
+
+        self.rope = LTX2RotaryPosEmbed1d(
+            self.inner_dim,
+            base_seq_len=rope_base_seq_len,
+            theta=rope_theta,
+            double_precision=rope_double_precision,
+            rope_type=rope_type,
+            num_attention_heads=num_attention_heads,
+        )
+
+        self.transformer_blocks = torch.nn.ModuleList(
+            [
+                LTX2TransformerBlock1d(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    rope_type=rope_type,
+                    apply_gated_attention=apply_gated_attention,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.norm_out = torch.nn.RMSNorm(
+            self.inner_dim, eps=eps, elementwise_affine=False
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        attn_mask_binarize_threshold: float = -9000.0,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # hidden_states shape: [batch_size, seq_len, hidden_dim]
+        # attention_mask shape: [batch_size, seq_len] or [batch_size, 1, 1, seq_len]
+        batch_size, seq_len, _ = hidden_states.shape
+
+        # 1. Replace padding with learned registers, if using
+        if self.learnable_registers is not None:
+            if seq_len % self.num_learnable_registers != 0:
+                raise ValueError(
+                    f"The `hidden_states` sequence length {hidden_states.shape[1]} should be divisible by the number"
+                    f" of learnable registers {self.num_learnable_registers}"
+                )
+
+            num_register_repeats = seq_len // self.num_learnable_registers
+            registers = torch.tile(
+                self.learnable_registers, (num_register_repeats, 1)
+            )  # [seq_len, inner_dim]
+
+            binary_attn_mask = (attention_mask >= attn_mask_binarize_threshold).int()
+            if binary_attn_mask.ndim == 4:
+                binary_attn_mask = binary_attn_mask.squeeze(1).squeeze(
+                    1
+                )  # [B, 1, 1, L] --> [B, L]
+
+            hidden_states_non_padded = [
+                hidden_states[i, binary_attn_mask[i].bool(), :]
+                for i in range(batch_size)
+            ]
+            valid_seq_lens = [x.shape[0] for x in hidden_states_non_padded]
+            pad_lengths = [seq_len - valid_seq_len for valid_seq_len in valid_seq_lens]
+            padded_hidden_states = [
+                F.pad(x, pad=(0, 0, 0, p), value=0)
+                for x, p in zip(hidden_states_non_padded, pad_lengths)
+            ]
+            padded_hidden_states = torch.cat(
+                [x.unsqueeze(0) for x in padded_hidden_states], dim=0
+            )  # [B, L, D]
+
+            flipped_mask = torch.flip(binary_attn_mask, dims=[1]).unsqueeze(
+                -1
+            )  # [B, L, 1]
+            hidden_states = (
+                flipped_mask * padded_hidden_states + (1 - flipped_mask) * registers
+            )
+
+            # Overwrite attention_mask with an all-zeros mask if using registers.
+            attention_mask = torch.zeros_like(attention_mask)
+
+        # 2. Calculate 1D RoPE positional embeddings
+        rotary_emb = self.rope(
+            batch_size,
+            seq_len,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        # 3. Run 1D transformer blocks
+        for block in self.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, attention_mask, rotary_emb
+                )
+            else:
+                hidden_states = block(
+                    hidden_states, attention_mask=attention_mask, rotary_emb=rotary_emb
+                )
+
+        hidden_states = self.norm_out(hidden_states)
+
+        return hidden_states, attention_mask
+
+
+class LTX2TextConnectors(nn.Module):
+    """
+    Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio
+    streams.
+    """
+
+    def __init__(
+        self,
+        config: LTX2ConnectorConfig,
+    ):
+        super().__init__()
+        caption_channels = config.caption_channels
+        self.caption_channels = caption_channels
+        text_proj_in_factor = config.text_proj_in_factor
+        video_connector_num_attention_heads = config.video_connector_num_attention_heads
+        video_connector_attention_head_dim = config.video_connector_attention_head_dim
+        video_connector_num_layers = config.video_connector_num_layers
+        video_connector_num_learnable_registers = (
+            config.video_connector_num_learnable_registers
+        )
+        audio_connector_num_attention_heads = config.audio_connector_num_attention_heads
+        audio_connector_attention_head_dim = config.audio_connector_attention_head_dim
+        audio_connector_num_layers = config.audio_connector_num_layers
+        audio_connector_num_learnable_registers = (
+            config.audio_connector_num_learnable_registers
+        )
+        connector_rope_base_seq_len = config.connector_rope_base_seq_len
+        rope_theta = config.rope_theta
+        rope_double_precision = config.rope_double_precision
+        causal_temporal_positioning = config.causal_temporal_positioning
+        rope_type = config.rope_type
+        connector_apply_gated_attention = config.connector_apply_gated_attention
+        feature_extractor_in_features = config.feature_extractor_in_features
+        video_feature_extractor_out_features = (
+            config.video_feature_extractor_out_features
+        )
+        audio_feature_extractor_out_features = (
+            config.audio_feature_extractor_out_features
+        )
+
+        self.text_proj_in: nn.Linear | None = None
+        self.video_aggregate_embed: nn.Linear | None = None
+        self.audio_aggregate_embed: nn.Linear | None = None
+        if (
+            feature_extractor_in_features > 0
+            and video_feature_extractor_out_features > 0
+            and audio_feature_extractor_out_features > 0
+        ):
+            self.video_aggregate_embed = nn.Linear(
+                feature_extractor_in_features,
+                video_feature_extractor_out_features,
+                bias=True,
+            )
+            self.audio_aggregate_embed = nn.Linear(
+                feature_extractor_in_features,
+                audio_feature_extractor_out_features,
+                bias=True,
+            )
+        else:
+            self.text_proj_in = nn.Linear(
+                caption_channels * text_proj_in_factor, caption_channels, bias=False
+            )
+        self.video_connector = LTX2ConnectorTransformer1d(
+            num_attention_heads=video_connector_num_attention_heads,
+            attention_head_dim=video_connector_attention_head_dim,
+            num_layers=video_connector_num_layers,
+            num_learnable_registers=video_connector_num_learnable_registers,
+            rope_base_seq_len=connector_rope_base_seq_len,
+            rope_theta=rope_theta,
+            rope_double_precision=rope_double_precision,
+            causal_temporal_positioning=causal_temporal_positioning,
+            rope_type=rope_type,
+            apply_gated_attention=connector_apply_gated_attention,
+        )
+        self.audio_connector = LTX2ConnectorTransformer1d(
+            num_attention_heads=audio_connector_num_attention_heads,
+            attention_head_dim=audio_connector_attention_head_dim,
+            num_layers=audio_connector_num_layers,
+            num_learnable_registers=audio_connector_num_learnable_registers,
+            rope_base_seq_len=connector_rope_base_seq_len,
+            rope_theta=rope_theta,
+            rope_double_precision=rope_double_precision,
+            causal_temporal_positioning=causal_temporal_positioning,
+            rope_type=rope_type,
+            apply_gated_attention=connector_apply_gated_attention,
+        )
+
+    @staticmethod
+    def _rescale_v2_features(
+        x: torch.Tensor, target_dim: int, source_dim: int
+    ) -> torch.Tensor:
+        return x * math.sqrt(target_dim / source_dim)
+
+    def forward(
+        self,
+        text_encoder_hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        additive_mask: bool = False,
+    ):
+        # Convert to additive attention mask, if necessary
+        if not additive_mask:
+            text_dtype = text_encoder_hidden_states.dtype
+            attention_mask = (attention_mask - 1).reshape(
+                attention_mask.shape[0], 1, -1, attention_mask.shape[-1]
+            )
+            attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max
+
+        # Ensure sequence length is divisible by num_learnable_registers (128)
+        seq_len = text_encoder_hidden_states.shape[1]
+        num_learnable_registers = self.video_connector.num_learnable_registers
+        if (
+            num_learnable_registers is not None
+            and seq_len % num_learnable_registers != 0
+        ):
+            pad_len = num_learnable_registers - (seq_len % num_learnable_registers)
+            text_encoder_hidden_states = F.pad(
+                text_encoder_hidden_states, (0, 0, 0, pad_len), value=0.0
+            )
+
+            if attention_mask.shape[-1] == seq_len:
+                # Pad with a large negative value to mask out the new tokens
+                attention_mask = F.pad(attention_mask, (0, pad_len), value=-1000000.0)
+
+        if (
+            self.video_aggregate_embed is not None
+            and self.audio_aggregate_embed is not None
+        ):
+            video_hidden_states = text_encoder_hidden_states
+            audio_hidden_states = text_encoder_hidden_states
+            if video_hidden_states.dtype != self.video_aggregate_embed.weight.dtype:
+                video_hidden_states = video_hidden_states.to(
+                    self.video_aggregate_embed.weight.dtype
+                )
+            if audio_hidden_states.dtype != self.audio_aggregate_embed.weight.dtype:
+                audio_hidden_states = audio_hidden_states.to(
+                    self.audio_aggregate_embed.weight.dtype
+                )
+            source_dim = self.caption_channels
+            video_hidden_states = self._rescale_v2_features(
+                video_hidden_states,
+                self.video_aggregate_embed.out_features,
+                source_dim,
+            )
+            audio_hidden_states = self._rescale_v2_features(
+                audio_hidden_states,
+                self.audio_aggregate_embed.out_features,
+                source_dim,
+            )
+            video_hidden_states = self.video_aggregate_embed(video_hidden_states)
+            audio_hidden_states = self.audio_aggregate_embed(audio_hidden_states)
+        else:
+            assert self.text_proj_in is not None
+            if text_encoder_hidden_states.dtype != self.text_proj_in.weight.dtype:
+                text_encoder_hidden_states = text_encoder_hidden_states.to(
+                    self.text_proj_in.weight.dtype
+                )
+            video_hidden_states = self.text_proj_in(text_encoder_hidden_states)
+            audio_hidden_states = video_hidden_states
+
+        video_text_embedding, new_attn_mask = self.video_connector(
+            video_hidden_states, attention_mask
+        )
+
+        attn_mask = (new_attn_mask < 1e-6).to(torch.int64)
+        attn_mask = attn_mask.reshape(
+            video_text_embedding.shape[0], video_text_embedding.shape[1], 1
+        )
+        video_text_embedding = video_text_embedding * attn_mask
+        new_attn_mask = attn_mask.squeeze(-1)
+
+        audio_text_embedding, _ = self.audio_connector(
+            audio_hidden_states, attention_mask
+        )
+
+        return video_text_embedding, audio_text_embedding, new_attn_mask
+
+
+EntryClass = LTX2TextConnectors

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,675 @@
+# 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.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+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,
+            # is_cross_attention=True,
+        )
+
+    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,
+    LayerwiseOffloadableModuleMixin,
+):
+    """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.
+    """
+
+    layerwise_offload_dit_group_enabled = False
+
+    _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_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/base.py

@@ -0,0 +1,127 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from abc import ABC, abstractmethod
+from typing import Any
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models import DiTConfig
+
+# NOTE: TeaCacheContext and TeaCacheMixin have been moved to
+# sglang.multimodal_gen.runtime.cache.teacache
+# For backwards compatibility, re-export from the new location
+from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheContext  # noqa: F401
+from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheMixin
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+# TODO
+class BaseDiT(nn.Module, ABC):
+    _fsdp_shard_conditions: list = []
+    _compile_conditions: list = []
+    param_names_mapping: dict
+    reverse_param_names_mapping: dict
+    hidden_size: int
+    num_attention_heads: int
+    num_channels_latents: int
+    # always supports torch_sdpa
+    _supported_attention_backends: set[AttentionBackendEnum] = (
+        DiTConfig()._supported_attention_backends
+    )
+
+    def __init_subclass__(cls) -> None:
+        required_class_attrs = [
+            "_fsdp_shard_conditions",
+            "param_names_mapping",
+            "_compile_conditions",
+        ]
+        super().__init_subclass__()
+        for attr in required_class_attrs:
+            if not hasattr(cls, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseDiT must define '{attr}' class variable"
+                )
+
+    def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        super().__init__()
+        self.config = config
+        self.hf_config = hf_config
+        if not self.supported_attention_backends:
+            raise ValueError(
+                f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
+            )
+
+    @abstractmethod
+    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,
+        **kwargs,
+    ) -> torch.Tensor:
+        pass
+
+    def __post_init__(self) -> None:
+        required_attrs = ["hidden_size", "num_attention_heads", "num_channels_latents"]
+        for attr in required_attrs:
+            if not hasattr(self, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseDiT must define '{attr}' instance variable"
+                )
+
+    def post_load_weights(self) -> None:
+        """Run model-specific post-load weight fixups after all parameters are materialized."""
+        return None
+
+    @property
+    def supported_attention_backends(self) -> set[AttentionBackendEnum]:
+        return self._supported_attention_backends
+
+    @property
+    def device(self) -> torch.device:
+        """Get the device of the model."""
+        return next(self.parameters()).device
+
+
+class CachableDiT(TeaCacheMixin, BaseDiT):
+    """
+    An intermediate base class that adds TeaCache optimization functionality to DiT models.
+
+    Inherits TeaCacheMixin for cache logic and BaseDiT for core DiT functionality.
+    """
+
+    # These are required class attributes that should be overridden by concrete implementations
+    _fsdp_shard_conditions = []
+    param_names_mapping = {}
+    reverse_param_names_mapping = {}
+    lora_param_names_mapping: dict = {}
+    # Ensure these instance attributes are properly defined in subclasses
+    hidden_size: int
+    num_attention_heads: int
+    num_channels_latents: int
+    # always supports torch_sdpa
+    _supported_attention_backends: set[AttentionBackendEnum] = (
+        DiTConfig()._supported_attention_backends
+    )
+
+    def __init__(self, config: DiTConfig, **kwargs) -> None:
+        super().__init__(config, **kwargs)
+        self._init_teacache_state()
+
+    @classmethod
+    def get_nunchaku_quant_rules(cls) -> dict[str, dict[str, Any]]:
+        """
+        Get quantization rules for Nunchaku quantization.
+
+        Returns a dict mapping layer name patterns to quantization configs:
+        {
+            "skip": [list of patterns to skip quantization],
+            "svdq_w4a4": [list of patterns for SVDQ W4A4],
+            "awq_w4a16": [list of patterns for AWQ W4A16],
+        }
+        """
+        return {}

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

@@ -0,0 +1,751 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch.nn.attention.flex_attention import (
+    BlockMask,
+    create_block_mask,
+    flex_attention,
+)
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+
+# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention
+# see https://github.com/pytorch/pytorch/issues/133254
+# change to default for other models
+flex_attention = torch.compile(
+    flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs"
+)
+import torch.distributed as dist
+
+from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_sp_world_size,
+    get_tp_rank,
+    get_tp_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
+from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
+    CausalSelfAttentionKVCache,
+    CrossAttentionKVCache,
+)
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    FP32LayerNorm,
+    LayerNormScaleShift,
+    RMSNorm,
+    ScaleResidualLayerNormScaleShift,
+    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.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
+    _apply_rotary_emb,
+    get_rotary_pos_embed,
+)
+from sglang.multimodal_gen.runtime.layers.visual_embedding import PatchEmbed
+from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT
+from sglang.multimodal_gen.runtime.models.dits.wanvideo import (
+    WanT2VCrossAttention,
+    WanTimeTextImageEmbedding,
+)
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils import add_prefix
+
+logger = init_logger(__name__)
+
+
+class CausalWanSelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        qk_norm=True,
+        eps=1e-6,
+        parallel_attention=False,
+        head_dim: int | None = None,
+        head_start: int = 0,
+    ) -> None:
+        if head_dim is None:
+            assert dim % num_heads == 0
+            head_dim = dim // num_heads
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.head_start = head_start
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.parallel_attention = parallel_attention
+
+        # Scaled dot product attention
+        self.attn = LocalAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=(
+                AttentionBackendEnum.FA,
+                AttentionBackendEnum.AITER,
+                AttentionBackendEnum.TORCH_SDPA,
+            ),
+        )
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        block_mask: BlockMask,
+        kv_cache: CausalSelfAttentionKVCache | None = None,
+        current_start: int = 0,
+        cache_start: int | None = None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        cos, sin = freqs_cis
+        roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
+        roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)
+
+        if kv_cache is None:
+            # Padding for flex attention
+            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+            padded_roped_query = torch.cat(
+                [
+                    roped_query,
+                    torch.zeros(
+                        [q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                        device=q.device,
+                        dtype=v.dtype,
+                    ),
+                ],
+                dim=1,
+            )
+
+            padded_roped_key = torch.cat(
+                [
+                    roped_key,
+                    torch.zeros(
+                        [k.shape[0], padded_length, k.shape[2], k.shape[3]],
+                        device=k.device,
+                        dtype=v.dtype,
+                    ),
+                ],
+                dim=1,
+            )
+
+            padded_v = torch.cat(
+                [
+                    v,
+                    torch.zeros(
+                        [v.shape[0], padded_length, v.shape[2], v.shape[3]],
+                        device=v.device,
+                        dtype=v.dtype,
+                    ),
+                ],
+                dim=1,
+            )
+
+            x = flex_attention(
+                query=padded_roped_query.transpose(2, 1),
+                key=padded_roped_key.transpose(2, 1),
+                value=padded_v.transpose(2, 1),
+                block_mask=block_mask,
+            )[:, :, :-padded_length].transpose(2, 1)
+        else:
+            if kv_cache.can_direct_current_attention(roped_key.shape[1]):
+                return self.attn(roped_query, roped_key, v)
+
+            cache_view = kv_cache.update_and_get_attention_kv(
+                key=roped_key,
+                value=v,
+                current_chunk_start=current_start,
+                cache_head_start=self.head_start,
+                debug_name="CausalWan KV cache",
+            )
+            x = self.attn(
+                roped_query,
+                cache_view.k,
+                cache_view.v,
+            )
+
+        return x
+
+
+class CausalWanTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        qk_norm: str = "rms_norm_across_heads",
+        cross_attn_norm: bool = False,
+        eps: float = 1e-6,
+        added_kv_proj_dim: int | None = None,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__()
+
+        # 1. Self-attention
+        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        use_megatron_tp = getattr(
+            self, "_use_megatron_tp", type(self) is CausalWanTransformerBlock
+        )
+        if use_megatron_tp:
+            self.to_q = ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=add_prefix("to_q", prefix),
+            )
+            self.to_k = ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=add_prefix("to_k", prefix),
+            )
+            self.to_v = ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=add_prefix("to_v", prefix),
+            )
+
+            self.to_out = RowParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                input_is_parallel=True,
+                quant_config=quant_config,
+                prefix=add_prefix("to_out", prefix),
+            )
+            # megatron-style tp shards the weight (qkv) column-wise, effectively splitting the attention heads
+            tp_size = get_tp_world_size()
+            self.local_num_heads = divide(num_heads, tp_size)
+            head_start = get_tp_rank() * self.local_num_heads
+        else:
+            self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
+            self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
+            self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
+            self.to_out = ReplicatedLinear(
+                dim, dim, bias=True, quant_config=quant_config
+            )
+            tp_size = 1
+            self.local_num_heads = num_heads
+            head_start = 0
+        dim_head = dim // num_heads
+        self.attn1 = CausalWanSelfAttention(
+            dim,
+            self.local_num_heads,
+            local_attn_size=local_attn_size,
+            sink_size=sink_size,
+            qk_norm=qk_norm,
+            eps=eps,
+            head_dim=dim_head,
+            head_start=head_start,
+        )
+        self.hidden_dim = dim
+        self.num_attention_heads = num_heads
+        self.local_attn_size = local_attn_size
+        self.dim_head = dim_head
+        if qk_norm == "rms_norm":
+            self.norm_q = RMSNorm(dim_head, eps=eps)
+            self.norm_k = RMSNorm(dim_head, eps=eps)
+        elif qk_norm == "rms_norm_across_heads":
+            # LTX applies qk norm across all heads
+            self.norm_q = RMSNorm(dim, eps=eps)
+            self.norm_k = RMSNorm(dim, eps=eps)
+        else:
+            print("QK Norm type not supported")
+            raise Exception
+        self.tp_rmsnorm = (
+            use_megatron_tp and qk_norm == "rms_norm_across_heads" and tp_size > 1
+        )
+        assert cross_attn_norm is True
+        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=True, dtype=torch.float32
+        )
+
+        # 2. Cross-attention
+        # Only T2V for now
+        cross_attn_backends = {
+            b for b in supported_attention_backends if not b.is_sparse
+        }
+        self.attn2 = WanT2VCrossAttention(
+            dim,
+            num_heads,
+            qk_norm=qk_norm,
+            eps=eps,
+            supported_attention_backends=cross_attn_backends,
+            quant_config=quant_config,
+        )
+        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
+        )
+
+        # 3. Feed-forward
+        self.ffn = MLP(
+            dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config
+        )
+        self.mlp_residual = MulAdd()
+
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        block_mask: BlockMask,
+        kv_cache: CausalSelfAttentionKVCache | None = None,
+        crossattn_cache: CrossAttentionKVCache | None = None,
+        current_start: int = 0,
+        cache_start: int | None = None,
+    ) -> torch.Tensor:
+        # hidden_states.shape: [batch_size, seq_length, inner_dim]
+        # temb.shape: [batch_size, num_frames, 6, inner_dim]
+        if hidden_states.dim() == 4:
+            hidden_states = hidden_states.squeeze(1)
+        num_frames = temb.shape[1]
+        frame_seqlen = hidden_states.shape[1] // num_frames
+        bs, seq_length, _ = hidden_states.shape
+        orig_dtype = hidden_states.dtype
+        # assert orig_dtype != torch.float32
+        e = self.scale_shift_table + temb.float()
+        # e.shape: [batch_size, num_frames, 6, inner_dim]
+        assert e.shape == (bs, num_frames, 6, self.hidden_dim)
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
+            6, dim=2
+        )
+        # *_msa.shape: [batch_size, num_frames, 1, inner_dim]
+        assert shift_msa.dtype == torch.float32
+
+        # 1. Self-attention
+        norm_hidden_states = (
+            (
+                self.norm1(hidden_states.float()).unflatten(
+                    dim=1, sizes=(num_frames, frame_seqlen)
+                )
+                * (1 + scale_msa)
+                + shift_msa
+            )
+            .flatten(1, 2)
+            .to(orig_dtype)
+        )
+        query, _ = self.to_q(norm_hidden_states)
+        key, _ = self.to_k(norm_hidden_states)
+        value, _ = self.to_v(norm_hidden_states)
+
+        if self.norm_q is not None:
+            if self.tp_rmsnorm:
+                query = tensor_parallel_rms_norm(query, self.norm_q)
+            else:
+                query = self.norm_q(query)
+        if self.norm_k is not None:
+            if self.tp_rmsnorm:
+                key = tensor_parallel_rms_norm(key, self.norm_k)
+            else:
+                key = self.norm_k(key)
+
+        query = query.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
+        key = key.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
+        value = value.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
+
+        attn_output = self.attn1(
+            query,
+            key,
+            value,
+            freqs_cis,
+            block_mask,
+            kv_cache,
+            current_start,
+            cache_start,
+        )
+        attn_output = attn_output.flatten(2)
+        attn_output, _ = self.to_out(attn_output)
+        attn_output = attn_output.squeeze(1)
+
+        null_shift = null_scale = torch.zeros(
+            (1,), device=hidden_states.device, dtype=hidden_states.dtype
+        )
+        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
+            hidden_states, attn_output, gate_msa, null_shift, null_scale
+        )
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype
+        ), hidden_states.to(orig_dtype)
+
+        # 2. Cross-attention
+        attn_output = self.attn2(
+            norm_hidden_states,
+            context=encoder_hidden_states,
+            context_lens=None,
+            crossattn_cache=crossattn_cache,
+        )
+        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
+            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa
+        )
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype
+        ), hidden_states.to(orig_dtype)
+
+        # 3. Feed-forward
+        ff_output = self.ffn(norm_hidden_states)
+        hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states)
+        hidden_states = hidden_states.to(orig_dtype)
+
+        return hidden_states
+
+
+class CausalWanTransformer3DModel(BaseDiT, LayerwiseOffloadableModuleMixin):
+    _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
+    _compile_conditions = WanVideoConfig()._compile_conditions
+    _supported_attention_backends = WanVideoConfig()._supported_attention_backends
+    param_names_mapping = WanVideoConfig().param_names_mapping
+    reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
+    lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
+
+    def __init__(
+        self,
+        config: WanVideoConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_dim = config.attention_head_dim
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_channels_latents = config.num_channels_latents
+        self.patch_size = config.patch_size
+        self.text_len = config.text_len
+        self.local_attn_size = config.local_attn_size
+
+        # 1. Patch & position embedding
+        self.patch_embedding = PatchEmbed(
+            in_chans=config.in_channels,
+            embed_dim=inner_dim,
+            patch_size=config.patch_size,
+            flatten=False,
+        )
+
+        # 2. Condition embeddings
+        self.condition_embedder = WanTimeTextImageEmbedding(
+            dim=inner_dim,
+            time_freq_dim=config.freq_dim,
+            text_embed_dim=config.text_dim,
+            image_embed_dim=config.image_dim,
+        )
+
+        # 3. Transformer blocks
+        self.blocks = nn.ModuleList(
+            [
+                CausalWanTransformerBlock(
+                    inner_dim,
+                    config.ffn_dim,
+                    config.num_attention_heads,
+                    config.local_attn_size,
+                    config.sink_size,
+                    config.qk_norm,
+                    config.cross_attn_norm,
+                    config.eps,
+                    config.added_kv_proj_dim,
+                    self._supported_attention_backends,
+                    prefix=f"{config.prefix}.blocks.{i}",
+                    quant_config=quant_config,
+                )
+                for i in range(config.num_layers)
+            ]
+        )
+
+        # 4. Output norm & projection
+        self.norm_out = LayerNormScaleShift(
+            inner_dim,
+            eps=config.eps,
+            elementwise_affine=False,
+            dtype=torch.float32,
+        )
+        self.proj_out = nn.Linear(
+            inner_dim, config.out_channels * math.prod(config.patch_size)
+        )
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(1, 2, inner_dim) / inner_dim**0.5
+        )
+
+        self.gradient_checkpointing = False
+
+        # Causal-specific
+        self.block_mask = None
+        self.num_frame_per_block = config.arch_config.num_frames_per_block
+        assert self.num_frame_per_block <= 3
+        self.independent_first_frame = False
+
+        self.__post_init__()
+
+        self.layer_names = [
+            "blocks",
+        ]
+
+    @staticmethod
+    def _prepare_blockwise_causal_attn_mask(
+        device: torch.device | str,
+        num_frames: int = 21,
+        frame_seqlen: int = 1560,
+        num_frame_per_block=1,
+        local_attn_size=-1,
+    ) -> BlockMask:
+        """
+        we will divide the token sequence into the following format
+        [1 latent frame] [1 latent frame] ... [1 latent frame]
+        We use flexattention to construct the attention mask
+        """
+        total_length = num_frames * frame_seqlen
+
+        # we do right padding to get to a multiple of 128
+        padded_length = math.ceil(total_length / 128) * 128 - total_length
+
+        ends = torch.zeros(
+            total_length + padded_length, device=device, dtype=torch.long
+        )
+
+        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+        frame_indices = torch.arange(
+            start=0,
+            end=total_length,
+            step=frame_seqlen * num_frame_per_block,
+            device=device,
+        )
+
+        for tmp in frame_indices:
+            ends[tmp : tmp + frame_seqlen * num_frame_per_block] = (
+                tmp + frame_seqlen * num_frame_per_block
+            )
+
+        def attention_mask(b, h, q_idx, kv_idx):
+            if local_attn_size == -1:
+                return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)
+            else:
+                return (
+                    (kv_idx < ends[q_idx])
+                    & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
+                ) | (q_idx == kv_idx)
+            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional mask
+
+        block_mask = create_block_mask(
+            attention_mask,
+            B=None,
+            H=None,
+            Q_LEN=total_length + padded_length,
+            KV_LEN=total_length + padded_length,
+            _compile=False,
+            device=device,
+        )
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
+            )
+            print(block_mask)
+
+        # import imageio
+        # import numpy as np
+        # from torch.nn.attention.flex_attention import create_mask
+
+        # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +
+        #                    padded_length, KV_LEN=total_length + padded_length, device=device)
+        # import cv2
+        # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))
+        # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))
+
+        return block_mask
+
+    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,
+        kv_cache: list[CausalSelfAttentionKVCache] | None = None,
+        crossattn_cache: list[CrossAttentionKVCache] | None = None,
+        current_start: int = 0,
+        cache_start: int = 0,
+        start_frame: int = 0,
+    ) -> torch.Tensor:
+        r"""
+        Run the diffusion model with kv caching.
+        See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.
+        This function will be run for num_frame times.
+        Process the latent frames one by one (1560 tokens each)
+        """
+
+        orig_dtype = hidden_states.dtype
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+        if (
+            isinstance(encoder_hidden_states_image, list)
+            and len(encoder_hidden_states_image) > 0
+        ):
+            encoder_hidden_states_image = encoder_hidden_states_image[0]
+        else:
+            encoder_hidden_states_image = None
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Get rotary embeddings
+        d = self.hidden_size // self.num_attention_heads
+        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (
+                post_patch_num_frames * get_sp_world_size(),
+                post_patch_height,
+                post_patch_width,
+            ),
+            self.hidden_size,
+            self.num_attention_heads,
+            rope_dim_list,
+            dtype=(
+                torch.float64
+                if current_platform.is_float64_supported()
+                else torch.float32
+            ),
+            rope_theta=10000,
+            start_frame=start_frame,  # Assume that start_frame is 0 when kv_cache is None
+        )
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        freqs_cis = (
+            (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None
+        )
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        (
+            temb,
+            timestep_proj,
+            encoder_hidden_states,
+            encoder_hidden_states_image,
+        ) = self.condition_embedder(
+            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image
+        )
+        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(
+            dim=0, sizes=timestep.shape
+        )
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat(
+                [encoder_hidden_states_image, encoder_hidden_states], dim=1
+            )
+
+        encoder_hidden_states = (
+            encoder_hidden_states.to(orig_dtype)
+            if current_platform.is_mps()
+            else encoder_hidden_states
+        )  # cast to orig_dtype for MPS
+
+        assert encoder_hidden_states.dtype == orig_dtype
+
+        # 4. Transformer blocks
+        for block_index, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                causal_kwargs = {
+                    "kv_cache": kv_cache[block_index],
+                    "current_start": current_start,
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask,
+                }
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    timestep_proj,
+                    freqs_cis,
+                    **causal_kwargs,
+                )
+            else:
+                causal_kwargs = {
+                    "kv_cache": kv_cache[block_index],
+                    "crossattn_cache": crossattn_cache[block_index],
+                    "current_start": current_start,
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask,
+                }
+                hidden_states = block(
+                    hidden_states,
+                    encoder_hidden_states,
+                    timestep_proj,
+                    freqs_cis,
+                    **causal_kwargs,
+                )
+
+        # 5. Output norm, projection & unpatchify
+        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
+        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2)
+        hidden_states = self.norm_out(hidden_states, shift, scale)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            post_patch_num_frames,
+            post_patch_height,
+            post_patch_width,
+            p_t,
+            p_h,
+            p_w,
+            -1,
+        )
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return output
+
+
+EntryClass = CausalWanTransformer3DModel

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

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+
+def modulate(
+    x: torch.Tensor,
+    shift: torch.Tensor | None = None,
+    scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Modulate by shift and scale."""
+    if scale is None and shift is None:
+        return x
+    if shift is None:
+        return x * (1 + scale.unsqueeze(1))  # type: ignore[union-attr]
+    if scale is None:
+        return x + shift.unsqueeze(1)  # type: ignore[union-attr]
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

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

@@ -0,0 +1,506 @@
+# Copyright 2026 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+
+from sglang.multimodal_gen.configs.models.dits.ernie_image import (
+    ErnieImageDitConfig,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+    get_tp_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.attention.layer import (
+    USPAttention,
+    build_varlen_mask_meta,
+)
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+
+
+def _rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)  # codespell:ignore nd
+    return out.float()
+
+
+class EmbedND3(nn.Module):
+    """3D rotary positional embedding for (temporal/batch_idx, height, width)."""
+
+    def __init__(self, dim: int, theta: int, axes_dim: Tuple[int, int, int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = list(axes_dim)
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        emb = torch.cat(
+            [_rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(3)],
+            dim=-1,
+        )
+        emb = emb.unsqueeze(1).permute(2, 0, 1, 3)
+        return torch.stack([emb, emb], dim=-1).reshape(*emb.shape[:-1], -1)
+
+
+class ErnieImageSelfAttention(nn.Module):
+    """Self-attention with separate Q/K/V projections and QK LayerNorm.
+
+    Module name hierarchy matches diffusers Attention naming convention:
+      self_attention.to_q, self_attention.to_k, self_attention.to_v,
+      self_attention.to_out.0, self_attention.norm_q, self_attention.norm_k.
+
+    Supports tensor parallelism: Q/K/V projections use ColumnParallelLinear
+    (output dim sharded by heads), output projection uses RowParallelLinear
+    (input dim sharded, all-reduce after matmul).
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        head_dim: int,
+        eps: float = 1e-6,
+        qk_layernorm: bool = True,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        tp_size = get_tp_world_size()
+        self.num_local_heads = num_heads // tp_size
+        assert (
+            num_heads % tp_size == 0
+        ), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})"
+
+        self.to_q = ColumnParallelLinear(
+            hidden_size,
+            hidden_size,
+            bias=False,
+            gather_output=False,
+            prefix=f"{prefix}.to_q",
+        )
+        self.to_k = ColumnParallelLinear(
+            hidden_size,
+            hidden_size,
+            bias=False,
+            gather_output=False,
+            prefix=f"{prefix}.to_k",
+        )
+        self.to_v = ColumnParallelLinear(
+            hidden_size,
+            hidden_size,
+            bias=False,
+            gather_output=False,
+            prefix=f"{prefix}.to_v",
+        )
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(
+                    hidden_size,
+                    hidden_size,
+                    bias=False,
+                    input_is_parallel=True,
+                    prefix=f"{prefix}.to_out.0",
+                ),
+            ]
+        )
+
+        self.qk_layernorm = qk_layernorm
+        if qk_layernorm:
+            self.norm_q = RMSNorm(head_dim, eps=eps)
+            self.norm_k = RMSNorm(head_dim, eps=eps)
+
+        # The joint [image, text] stream is fully replicated, so the ulysses
+        # all-to-all would wrongly treat it as sharded and duplicate it. Skip
+        # SP until the stream is sharded (sp_shard + num_replicated_suffix).
+        self.attn = USPAttention(
+            num_heads=self.num_local_heads,
+            head_size=head_dim,
+            prefix=f"{prefix}.attn",
+            skip_sequence_parallel=True,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rotary_pos_emb: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        attn_mask_meta: dict | None = None,
+    ) -> torch.Tensor:
+        B, S, H = x.shape
+
+        q, _ = self.to_q(x)
+        k, _ = self.to_k(x)
+        v, _ = self.to_v(x)
+
+        q = q.view(B, S, self.num_local_heads, self.head_dim)
+        k = k.view(B, S, self.num_local_heads, self.head_dim)
+        v = v.view(B, S, self.num_local_heads, self.head_dim)
+
+        if self.qk_layernorm:
+            q, k = apply_qk_norm(
+                q,
+                k,
+                self.norm_q,
+                self.norm_k,
+                self.head_dim,
+            )
+
+        q = _apply_rotary_bshd(q, rotary_pos_emb)
+        k = _apply_rotary_bshd(k, rotary_pos_emb)
+
+        attn_out = self.attn(
+            q, k, v, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta
+        )
+        attn_out = attn_out.reshape(B, S, self.num_local_heads * self.head_dim)
+        out, _ = self.to_out[0](attn_out)
+        return out
+
+
+class ErnieImageMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        ffn_hidden_size: int,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [ffn_hidden_size, ffn_hidden_size],
+            bias=False,
+            gather_output=False,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.linear_fc2 = RowParallelLinear(
+            ffn_hidden_size,
+            hidden_size,
+            bias=False,
+            input_is_parallel=True,
+            prefix=f"{prefix}.linear_fc2",
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        gate, up = gate_up.chunk(2, dim=-1)
+        x = up * F.gelu(gate)
+        x, _ = self.linear_fc2(x)
+        return x
+
+
+class ErnieImageSharedAdaLNBlock(nn.Module):
+    """Single-stream transformer block with externally-computed Shared AdaLN."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        head_dim: int,
+        ffn_hidden_size: int,
+        eps: float = 1e-6,
+        qk_layernorm: bool = True,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.adaLN_sa_ln = RMSNorm(hidden_size, eps=eps)
+        self.self_attention = ErnieImageSelfAttention(
+            hidden_size,
+            num_heads,
+            head_dim,
+            eps,
+            qk_layernorm,
+            prefix=f"{prefix}.self_attention",
+        )
+        self.adaLN_mlp_ln = RMSNorm(hidden_size, eps=eps)
+        self.mlp = ErnieImageMLP(hidden_size, ffn_hidden_size, prefix=f"{prefix}.mlp")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rotary_pos_emb: torch.Tensor,
+        shift_msa: torch.Tensor,
+        scale_msa: torch.Tensor,
+        gate_msa: torch.Tensor,
+        shift_mlp: torch.Tensor,
+        scale_mlp: torch.Tensor,
+        gate_mlp: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        attn_mask_meta: dict | None = None,
+    ) -> torch.Tensor:
+        residual = x
+        x = self.adaLN_sa_ln(x) * (1 + scale_msa) + shift_msa
+        x = residual + gate_msa * self.self_attention(
+            x, rotary_pos_emb, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta
+        )
+
+        residual = x
+        x = self.adaLN_mlp_ln(x) * (1 + scale_mlp) + shift_mlp
+        x = residual + gate_mlp * self.mlp(x)
+
+        return x
+
+
+def _apply_rotary_bshd(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    freqs = freqs.permute(1, 0, 2, 3)
+    rot_dim = freqs.shape[-1]
+    x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
+
+    cos_ = torch.cos(freqs).to(x.dtype)
+    sin_ = torch.sin(freqs).to(x.dtype)
+
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    x_rotated = torch.cat((-x2, x1), dim=-1)
+
+    x_rot = x_rot * cos_ + x_rotated * sin_
+    return torch.cat((x_rot, x_pass), dim=-1)
+
+
+class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """ErnieImage DiT: Single-stream transformer with Shared AdaLN."""
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["ErnieImageSharedAdaLNBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+
+    _fsdp_shard_conditions = ErnieImageDitConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = []
+    param_names_mapping = ErnieImageDitConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = {}
+
+    def __init__(
+        self,
+        config: ErnieImageDitConfig,
+        hf_config: dict[str, Any],
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+
+        arch = config.arch_config
+        self.hidden_size = arch.hidden_size
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.out_channels
+        self.head_dim = arch.attention_head_dim
+        self.num_layers = arch.num_layers
+        self.patch_size = arch.patch_size
+        self.out_channels = arch.out_channels
+        self.inner_dim = self.hidden_size
+
+        self.x_embedder = nn.ModuleDict(
+            {
+                "proj": nn.Conv2d(
+                    arch.in_channels,
+                    self.inner_dim,
+                    kernel_size=arch.patch_size,
+                    stride=arch.patch_size,
+                    bias=True,
+                ),
+            }
+        )
+
+        if arch.text_in_dim != self.inner_dim:
+            self.text_proj = nn.Linear(arch.text_in_dim, self.inner_dim, bias=False)
+        else:
+            self.text_proj = None
+
+        self.time_proj = Timesteps(
+            self.inner_dim,
+            flip_sin_to_cos=False,
+            downscale_freq_shift=0,
+        )
+        self.time_embedding = TimestepEmbedding(
+            in_channels=self.inner_dim,
+            time_embed_dim=self.inner_dim,
+        )
+
+        self.pos_embed = EmbedND3(
+            dim=self.head_dim,
+            theta=arch.rope_theta,
+            axes_dim=arch.rope_axes_dim,
+        )
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(self.inner_dim, 6 * self.inner_dim),
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                ErnieImageSharedAdaLNBlock(
+                    hidden_size=self.inner_dim,
+                    num_heads=self.num_attention_heads,
+                    head_dim=self.head_dim,
+                    ffn_hidden_size=arch.ffn_hidden_size,
+                    eps=arch.eps,
+                    qk_layernorm=arch.qk_layernorm,
+                    prefix=f"layers.{i}",
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+
+        self.final_norm = nn.ModuleDict(
+            {
+                "norm": nn.LayerNorm(
+                    self.inner_dim, elementwise_affine=False, eps=arch.eps
+                ),
+                "linear": nn.Linear(self.inner_dim, self.inner_dim * 2),
+            }
+        )
+
+        self.final_linear = ColumnParallelLinear(
+            self.inner_dim,
+            arch.patch_size * arch.patch_size * self.out_channels,
+            bias=True,
+            gather_output=True,
+            prefix="final_linear",
+        )
+
+        self.layer_names = ["layers"]
+
+        self.__post_init__()
+
+    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,
+        encoder_hidden_states_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [B, C, H, W] latent images (patchified, 128 channels)
+            encoder_hidden_states: [B, T, text_dim] or list of text embeddings
+            timestep: [B] timestep values
+        Returns:
+            output: [B, C, H, W] predicted noise / denoised output
+        """
+        device, dtype = hidden_states.device, hidden_states.dtype
+        B, C, H, W = hidden_states.shape
+        p = self.patch_size
+        Hp, Wp = H // p, W // p
+        N_img = Hp * Wp
+
+        img_tokens = self.x_embedder["proj"](hidden_states)  # [B, D, Hp, Wp]
+        img_tokens = img_tokens.reshape(B, self.inner_dim, N_img).transpose(
+            1, 2
+        )  # [B, N_img, D]
+
+        if isinstance(encoder_hidden_states, (list, tuple)):
+            encoder_hidden_states = encoder_hidden_states[0]
+        text_tokens = encoder_hidden_states  # [B, T, text_dim]
+        if self.text_proj is not None and text_tokens.numel() > 0:
+            text_tokens = self.text_proj(text_tokens)
+        Tmax = text_tokens.shape[1]
+
+        x = torch.cat([img_tokens, text_tokens], dim=1)  # [B, S, D]
+
+        grid_yx = torch.stack(
+            torch.meshgrid(
+                torch.arange(Hp, device=device, dtype=torch.float32),
+                torch.arange(Wp, device=device, dtype=torch.float32),
+                indexing="ij",
+            ),
+            dim=-1,
+        ).reshape(-1, 2)
+
+        image_ids = torch.cat(
+            [
+                torch.full((B, N_img, 1), Tmax, device=device, dtype=torch.float32),
+                grid_yx.view(1, N_img, 2).expand(B, -1, -1),
+            ],
+            dim=-1,
+        )
+
+        if Tmax > 0:
+            text_ids = torch.cat(
+                [
+                    torch.arange(Tmax, device=device, dtype=torch.float32)
+                    .view(1, Tmax, 1)
+                    .expand(B, -1, -1),
+                    torch.zeros((B, Tmax, 2), device=device),
+                ],
+                dim=-1,
+            )
+        else:
+            text_ids = torch.zeros((B, 0, 3), device=device)
+
+        all_ids = torch.cat([image_ids, text_ids], dim=1)
+        rotary_pos_emb = self.pos_embed(all_ids)
+
+        attn_mask = attn_mask_meta = None
+        if encoder_hidden_states_mask is not None:
+            image_mask = torch.ones((B, N_img), dtype=torch.bool, device=device)
+            attn_mask = torch.cat(
+                [
+                    image_mask,
+                    encoder_hidden_states_mask.to(device=device, dtype=torch.bool),
+                ],
+                dim=1,
+            )
+            attn_mask_meta = build_varlen_mask_meta(attn_mask)
+
+        t_emb = self.time_proj(timestep.to(dtype))
+        c = self.time_embedding(t_emb.to(dtype=dtype))
+
+        mod_params = self.adaLN_modulation(c)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            t.unsqueeze(1) for t in mod_params.chunk(6, dim=-1)
+        )
+
+        for layer in self.layers:
+            x = layer(
+                x,
+                rotary_pos_emb,
+                shift_msa,
+                scale_msa,
+                gate_msa,
+                shift_mlp,
+                scale_mlp,
+                gate_mlp,
+                attn_mask=attn_mask,
+                attn_mask_meta=attn_mask_meta,
+            )
+
+        scale, shift = self.final_norm["linear"](c).chunk(2, dim=-1)
+        x = self.final_norm["norm"](x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+        patches, _ = self.final_linear(x[:, :N_img, :])
+
+        output = patches.view(B, Hp, Wp, p, p, self.out_channels)
+        output = output.permute(0, 5, 1, 3, 2, 4).contiguous()
+        output = output.view(B, self.out_channels, H, W)
+
+        return output
+
+
+EntryClass = ErnieImageTransformer2DModel

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

@@ -0,0 +1,976 @@
+# Copyright 2025 The CogView team, Tsinghua University & ZhipuAI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_sp_parallel_rank,
+    get_sp_world_size,
+    get_tp_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    ScaleResidualLayerNormScaleShift,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.mlp import FeedForward
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
+    _apply_rotary_emb,
+    apply_flashinfer_rope_qk_inplace,
+)
+from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_is_cuda = current_platform.is_cuda()
+
+
+class GlmImageLayerKVCache:
+    """KV cache for GlmImage model."""
+
+    def __init__(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode: Optional[str] = None  # "write", "read", "skip"
+
+    def store(self, k: torch.Tensor, v: torch.Tensor):
+        if self.k_cache is None:
+            self.k_cache = k
+            self.v_cache = v
+        else:
+            self.k_cache = torch.cat([self.k_cache, k], dim=1)
+            self.v_cache = torch.cat([self.v_cache, v], dim=1)
+
+    def get(self):
+        return self.k_cache, self.v_cache
+
+    def clear(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode = None
+
+
+class GlmImageKVCache:
+    """Container for all layers' KV caches."""
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+        self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
+
+    def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
+        return self.caches[layer_idx]
+
+    def set_mode(self, mode: Optional[str]):
+        if mode is not None and mode not in ["write", "read", "skip"]:
+            raise ValueError(
+                f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'"
+            )
+        for cache in self.caches:
+            cache.mode = mode
+
+    def clear(self):
+        for cache in self.caches:
+            cache.clear()
+
+
+class GlmImageTimestepEmbedding(nn.Module):
+    """
+    Replacement for diffusers TimestepEmbedding using ReplicatedLinear.
+    Structure: linear_1 -> act(silu) -> linear_2
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+    ):
+        super().__init__()
+        if out_dim is None:
+            out_dim = time_embed_dim
+        self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True)
+        if act_fn == "silu":
+            self.act = nn.SiLU()
+        elif act_fn == "gelu":
+            self.act = nn.GELU(approximate="tanh")
+        else:
+            self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True)
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        sample, _ = self.linear_1(sample)
+        sample = self.act(sample)
+        sample, _ = self.linear_2(sample)
+        return sample
+
+
+class GlmImageTextProjection(nn.Module):
+    """
+    Replacement for diffusers PixArtAlphaTextProjection using ReplicatedLinear.
+    Structure: linear_1 -> act_1 -> linear_2
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_size: int,
+        out_features: int = None,
+        act_fn: str = "silu",
+    ):
+        super().__init__()
+        if out_features is None:
+            out_features = hidden_size
+        self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True)
+        if act_fn == "silu":
+            self.act_1 = nn.SiLU()
+        elif act_fn == "gelu_tanh":
+            self.act_1 = nn.GELU(approximate="tanh")
+        else:
+            self.act_1 = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True)
+
+    def forward(self, caption: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.linear_1(caption)
+        hidden_states = self.act_1(hidden_states)
+        hidden_states, _ = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class GlmImageCombinedTimestepSizeEmbeddings(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        condition_dim: int,
+        pooled_projection_dim: int,
+        timesteps_dim: int = 256,
+    ):
+        super().__init__()
+
+        self.time_proj = Timesteps(
+            num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.condition_proj = Timesteps(
+            num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.timestep_embedder = GlmImageTimestepEmbedding(
+            in_channels=timesteps_dim, time_embed_dim=embedding_dim
+        )
+        self.condition_embedder = GlmImageTextProjection(
+            pooled_projection_dim, embedding_dim, act_fn="silu"
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        hidden_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+
+        crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(
+            crop_coords.size(0), -1
+        )
+        target_size_proj = self.condition_proj(target_size.flatten()).view(
+            target_size.size(0), -1
+        )
+
+        # (B, 2 * condition_dim)
+        condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1)
+
+        timesteps_emb = self.timestep_embedder(
+            timesteps_proj.to(dtype=hidden_dtype)
+        )  # (B, embedding_dim)
+        condition_emb = self.condition_embedder(
+            condition_proj.to(dtype=hidden_dtype)
+        )  # (B, embedding_dim)
+
+        conditioning = timesteps_emb + condition_emb
+        return conditioning
+
+
+class GlmImageImageProjector(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 16,
+        hidden_size: int = 2560,
+        patch_size: int = 2,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+
+        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, channel, height, width = hidden_states.shape
+        post_patch_height = height // self.patch_size
+        post_patch_width = width // self.patch_size
+
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            channel,
+            post_patch_height,
+            self.patch_size,
+            post_patch_width,
+            self.patch_size,
+        )
+        hidden_states = (
+            hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
+        )
+        hidden_states = self.proj(hidden_states)
+
+        return hidden_states
+
+
+class GlmImageAdaLayerNormZero(nn.Module):
+    def __init__(self, embedding_dim: int, dim: int) -> None:
+        super().__init__()
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        dtype = hidden_states.dtype
+        norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
+        norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(
+            dtype=dtype
+        )
+
+        emb, _ = self.linear(temb)
+        (
+            shift_msa,
+            c_shift_msa,
+            scale_msa,
+            c_scale_msa,
+            gate_msa,
+            c_gate_msa,
+            shift_mlp,
+            c_shift_mlp,
+            scale_mlp,
+            c_scale_mlp,
+            gate_mlp,
+            c_gate_mlp,
+        ) = emb.chunk(12, dim=1)
+
+        hidden_states = norm_hidden_states * (
+            1 + scale_msa.unsqueeze(1)
+        ) + shift_msa.unsqueeze(1)
+        encoder_hidden_states = norm_encoder_hidden_states * (
+            1 + c_scale_msa.unsqueeze(1)
+        ) + c_shift_msa.unsqueeze(1)
+
+        return (
+            hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        )
+
+
+class GlmImageGELU(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inner_dim: int,
+        bias: bool = True,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.proj = ColumnParallelLinear(
+            dim,
+            inner_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj" if prefix else "proj",
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.proj(hidden_states)
+        return F.gelu(hidden_states, approximate="tanh")
+
+
+class GlmImageFeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        inner_dim: Optional[int] = None,
+        bias: bool = True,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if inner_dim is None:
+            inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        self.net = nn.ModuleList(
+            [
+                GlmImageGELU(
+                    dim,
+                    inner_dim,
+                    bias=bias,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.net.0" if prefix else "net.0",
+                ),
+                nn.Dropout(0.0),
+                RowParallelLinear(
+                    inner_dim,
+                    dim_out,
+                    bias=bias,
+                    input_is_parallel=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.net.2" if prefix else "net.2",
+                ),
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.net[0](hidden_states)
+        hidden_states = self.net[1](hidden_states)
+        hidden_states, _ = self.net[2](hidden_states)
+        return hidden_states
+
+
+class GlmImageAttention(torch.nn.Module):
+    def __init__(
+        self,
+        query_dim,
+        heads,
+        dim_head,
+        out_dim,
+        bias,
+        qk_norm,
+        elementwise_affine,
+        eps,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__()
+
+        self.k_cache = None
+        self.v_cache = None
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.dim_head = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        tp_size = get_tp_world_size()
+        assert (
+            self.heads % tp_size == 0
+        ), f"heads ({self.heads}) must be divisible by tp_size ({tp_size})"
+        self.num_local_heads = self.heads // tp_size
+        self.num_local_kv_heads = self.num_local_heads
+
+        self.to_q = ColumnParallelLinear(
+            query_dim,
+            self.inner_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_q" if prefix else "to_q",
+        )
+        self.to_k = ColumnParallelLinear(
+            query_dim,
+            self.inner_kv_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_k" if prefix else "to_k",
+        )
+        self.to_v = ColumnParallelLinear(
+            query_dim,
+            self.inner_kv_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_v" if prefix else "to_v",
+        )
+
+        # (dropout omitted)
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(
+                    self.inner_dim,
+                    self.out_dim,
+                    bias=True,
+                    input_is_parallel=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0" if prefix else "to_out.0",
+                )
+            ]
+        )
+
+        if qk_norm is None:
+            self.norm_q = None
+            self.norm_k = None
+        elif qk_norm == "layer_norm":
+            self.norm_q = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+            self.norm_k = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+        else:
+            raise ValueError(
+                f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'."
+            )
+
+        self.attn = USPAttention(
+            num_heads=self.num_local_heads,
+            head_size=dim_head,
+            num_kv_heads=self.num_local_kv_heads,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        dtype = encoder_hidden_states.dtype
+
+        batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
+        batch_size, image_seq_length, embed_dim = hidden_states.shape
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # 1. QKV projections
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(hidden_states)
+        value, _ = self.to_v(hidden_states)
+
+        query = query.unflatten(2, (self.num_local_heads, -1))
+        key = key.unflatten(2, (self.num_local_kv_heads, -1))
+        value = value.unflatten(2, (self.num_local_kv_heads, -1))
+
+        # 2. QK normalization
+        if self.norm_q is not None:
+            query = self.norm_q(query).to(dtype=dtype)
+        if self.norm_k is not None:
+            key = self.norm_k(key).to(dtype=dtype)
+
+        # 3. Rotational positional embeddings applied to latent stream
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+
+            if _is_cuda and cos.dim() == 2:
+                q_img = query[:, text_seq_length:, :, :]
+                k_img = key[:, text_seq_length:, :, :]
+                cos_sin_cache = torch.cat(
+                    [
+                        cos.to(dtype=torch.float32).contiguous(),
+                        sin.to(dtype=torch.float32).contiguous(),
+                    ],
+                    dim=-1,
+                )
+                # apply_flashinfer_rope_qk_inplace is inplace kernel and q_img/k_img are views of query/key, so we need not copy back
+                q_out, k_out = apply_flashinfer_rope_qk_inplace(
+                    q_img, k_img, cos_sin_cache, is_neox=True
+                )
+            else:
+                query[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                    query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+                )
+                key[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                    key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+                )
+
+        if kv_cache is not None:
+            if kv_cache.mode == "write":
+                kv_cache.store(key, value)
+            elif kv_cache.mode == "read":
+                k_cache, v_cache = kv_cache.get()
+                key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key
+                value = (
+                    torch.cat([v_cache, value], dim=1) if v_cache is not None else value
+                )
+            elif kv_cache.mode == "skip":
+                pass
+
+        # 4. Attention
+        if attention_mask is not None:
+            text_attn_mask = attention_mask
+            assert (
+                text_attn_mask.dim() == 2
+            ), "the shape of text_attn_mask should be (batch_size, text_seq_length)"
+        hidden_states = self.attn(
+            query, key, value, num_replicated_prefix=text_seq_length
+        )
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # 5. Output projection
+        hidden_states, _ = self.to_out[0](hidden_states)
+        # hidden_states = self.to_out[1](hidden_states)         # (dropout omitted)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+class GlmImageTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int = 2560,
+        num_attention_heads: int = 64,
+        attention_head_dim: int = 40,
+        time_embed_dim: int = 512,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__()
+
+        # 1. Attention
+        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
+
+        self.attn1 = GlmImageAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            out_dim=dim,
+            bias=True,
+            qk_norm="layer_norm",
+            elementwise_affine=False,
+            eps=1e-5,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn1",
+            quant_config=quant_config,
+        )
+
+        # 2. Feedforward
+        self.norm2 = ScaleResidualLayerNormScaleShift(
+            dim, eps=1e-5, elementwise_affine=False
+        )
+        self.norm2_context = ScaleResidualLayerNormScaleShift(
+            dim, eps=1e-5, elementwise_affine=False
+        )
+        self.ff = GlmImageFeedForward(
+            dim=dim,
+            dim_out=dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff" if prefix else "ff",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        attention_mask: Optional[Dict[str, torch.Tensor]] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Timestep conditioning
+        (
+            norm_hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            norm_encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        ) = self.norm1(hidden_states, encoder_hidden_states, temb)
+
+        # 2. Attention
+        if attention_kwargs is None:
+            attention_kwargs = {}
+
+        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            attention_mask=attention_mask,
+            kv_cache=kv_cache,
+            **attention_kwargs,
+        )
+
+        # 3. Feedforward (fused residual + norm + scale/shift)
+        norm_hidden_states, hidden_states = self.norm2(
+            hidden_states,
+            attn_hidden_states,
+            gate_msa.unsqueeze(1),
+            shift_mlp.unsqueeze(1),
+            scale_mlp.unsqueeze(1),
+        )
+        norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context(
+            encoder_hidden_states,
+            attn_encoder_hidden_states,
+            c_gate_msa.unsqueeze(1),
+            c_shift_mlp.unsqueeze(1),
+            c_scale_mlp.unsqueeze(1),
+        )
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output_context = self.ff(norm_encoder_hidden_states)
+        hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
+        encoder_hidden_states = (
+            encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)
+        )
+
+        return hidden_states, encoder_hidden_states
+
+
+class GlmImageRotaryPosEmbed(nn.Module):
+    def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
+        super().__init__()
+
+        self.dim = dim
+        self.patch_size = patch_size
+        self.theta = theta
+
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_channels, height, width = hidden_states.shape
+        height, width = height // self.patch_size, width // self.patch_size
+        device = hidden_states.device
+
+        dim_h, dim_w = self.dim // 2, self.dim // 2
+        h_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[
+                    : (dim_h // 2)
+                ].float()
+                / dim_h
+            )
+        )
+        w_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[
+                    : (dim_w // 2)
+                ].float()
+                / dim_w
+            )
+        )
+        h_seq = torch.arange(height, device=device)
+        w_seq = torch.arange(width, device=device)
+        freqs_h = torch.outer(h_seq, h_inv_freq)
+        freqs_w = torch.outer(w_seq, w_inv_freq)
+
+        # Create position matrices for height and width
+        # [height, 1, dim//4] and [1, width, dim//4]
+        freqs_h = freqs_h.unsqueeze(1)
+        freqs_w = freqs_w.unsqueeze(0)
+        # Broadcast freqs_h and freqs_w to [height, width, dim//4]
+        freqs_h = freqs_h.expand(height, width, -1)
+        freqs_w = freqs_w.expand(height, width, -1)
+
+        # Concatenate along last dimension to get [height, width, dim//2]
+        freqs = torch.cat([freqs_h, freqs_w], dim=-1)
+        freqs = freqs.reshape(height * width, -1)  # [height * width, dim//2]
+        return (freqs.cos(), freqs.sin())
+
+
+class GlmImageAdaLayerNormContinuous(nn.Module):
+    """
+    GlmImage-only final AdaLN: LN(x) -> Linear(cond) -> chunk -> affine. Matches Megatron: **no activation** before the
+    Linear on conditioning embedding.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine: bool = True,
+        eps: float = 1e-5,
+        bias: bool = True,
+        norm_type: str = "layer_norm",
+    ):
+        super().__init__()
+        self.linear = nn.Linear(
+            conditioning_embedding_dim, embedding_dim * 2, bias=bias
+        )
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+            # For now, don’t replace this with sglang’s LayerNorm
+            # because the model doesn’t have this parameter and it will break model loading
+        elif norm_type == "rms_norm":
+            self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self, x: torch.Tensor, conditioning_embedding: torch.Tensor
+    ) -> torch.Tensor:
+        # *** NO SiLU here ***
+        emb = self.linear(conditioning_embedding.to(x.dtype))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+class GlmImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    r"""
+    Args:
+        patch_size (`int`, defaults to `2`):
+            The size of the patches to use in the patch embedding layer.
+        in_channels (`int`, defaults to `16`):
+            The number of channels in the input.
+        num_layers (`int`, defaults to `30`):
+            The number of layers of Transformer blocks to use.
+        attention_head_dim (`int`, defaults to `40`):
+            The number of channels in each head.
+        num_attention_heads (`int`, defaults to `64`):
+            The number of heads to use for multi-head attention.
+        out_channels (`int`, defaults to `16`):
+            The number of channels in the output.
+        text_embed_dim (`int`, defaults to `1472`):
+            Input dimension of text embeddings from the text encoder.
+        time_embed_dim (`int`, defaults to `512`):
+            Output dimension of timestep embeddings.
+        condition_dim (`int`, defaults to `256`):
+            The embedding dimension of the input SDXL-style resolution conditions (original_size, target_size,
+            crop_coords).
+        pos_embed_max_size (`int`, defaults to `128`):
+            The maximum resolution of the positional embeddings, from which slices of shape `H x W` are taken and added
+            to input patched latents, where `H` and `W` are the latent height and width respectively. A value of 128
+            means that the maximum supported height and width for image generation is `128 * vae_scale_factor *
+            patch_size => 128 * 8 * 2 => 2048`.
+        sample_size (`int`, defaults to `128`):
+            The base resolution of input latents. If height/width is not provided during generation, this value is used
+            to determine the resolution as `sample_size * vae_scale_factor => 128 * 8 => 1024`
+    """
+
+    def __init__(
+        self,
+        config: GlmImageDitConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+
+        self.config_data = config  # Store config
+        arch_config = config.arch_config
+
+        self.in_channels = arch_config.in_channels
+        self.out_channels = arch_config.out_channels
+        self.patch_size = arch_config.patch_size
+        self.num_layers = arch_config.num_layers
+        self.attention_head_dim = arch_config.attention_head_dim
+        self.num_attention_heads = arch_config.num_attention_heads
+        self.text_embed_dim = arch_config.text_embed_dim
+        self.time_embed_dim = arch_config.time_embed_dim
+
+        # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
+        # Each of these are sincos embeddings of shape 2 * condition_dim
+        pooled_projection_dim = 2 * 2 * arch_config.condition_dim
+        inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim
+
+        # 1. RoPE
+        self.rotary_emb = GlmImageRotaryPosEmbed(
+            arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0
+        )
+
+        # 2. Patch & Text-timestep embedding
+        self.image_projector = GlmImageImageProjector(
+            arch_config.in_channels, inner_dim, arch_config.patch_size
+        )
+        self.glyph_projector = FeedForward(
+            arch_config.text_embed_dim,
+            inner_dim,
+            inner_dim=inner_dim,
+            activation_fn="gelu",
+        )
+        self.prior_token_embedding = nn.Embedding(
+            arch_config.prior_vq_quantizer_codebook_size, inner_dim
+        )
+        self.prior_projector = FeedForward(
+            inner_dim, inner_dim, inner_dim=inner_dim, activation_fn="linear-silu"
+        )
+
+        self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
+            embedding_dim=arch_config.time_embed_dim,
+            condition_dim=arch_config.condition_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            timesteps_dim=arch_config.time_embed_dim,
+        )
+
+        # 3. Transformer blocks
+        self._supported_attention_backends = arch_config._supported_attention_backends
+        self.transformer_blocks = nn.ModuleList(
+            [
+                GlmImageTransformerBlock(
+                    inner_dim,
+                    arch_config.num_attention_heads,
+                    arch_config.attention_head_dim,
+                    arch_config.time_embed_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                    prefix=f"transformer_blocks.{i}",
+                    quant_config=quant_config,
+                )
+                for i in range(arch_config.num_layers)
+            ]
+        )
+
+        # 4. Output projection
+        self.norm_out = GlmImageAdaLayerNormContinuous(
+            inner_dim, arch_config.time_embed_dim, elementwise_affine=False
+        )
+        self.proj_out = nn.Linear(
+            inner_dim,
+            arch_config.patch_size * arch_config.patch_size * arch_config.out_channels,
+            bias=True,
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        prior_token_id: torch.Tensor,
+        prior_token_drop: torch.Tensor,
+        timestep: torch.LongTensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        kv_caches: Optional[GlmImageKVCache] = None,
+        kv_caches_mode: Optional[str] = None,
+        freqs_cis: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        ###
+        guidance: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor]:
+        if kv_caches is not None:
+            kv_caches.set_mode(kv_caches_mode)
+
+        batch_size, num_channels, height, width = hidden_states.shape
+
+        timestep = timestep - 1.0
+
+        if isinstance(encoder_hidden_states, list):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        # 1. RoPE
+        image_rotary_emb = freqs_cis
+        if image_rotary_emb is None:
+            image_rotary_emb = self.rotary_emb(hidden_states)
+        # 2. Patch & Timestep embeddings
+        p = self.config.patch_size
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        hidden_states = self.image_projector(hidden_states)
+        encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
+        prior_embedding = self.prior_token_embedding(prior_token_id)
+        prior_embedding = prior_embedding.masked_fill(prior_token_drop.unsqueeze(-1), 0)
+        prior_hidden_states = self.prior_projector(prior_embedding)
+        # SP: when latents are H-sharded, hidden_states has fewer patches than prior_hidden_states.
+        # Shard prior_hidden_states along seq dim to match (prior is row-major, same as latent patches).
+        if (
+            get_sp_world_size() > 1
+            and prior_hidden_states.shape[1] != hidden_states.shape[1]
+        ):
+            rank = get_sp_parallel_rank()
+            sp_world_size = get_sp_world_size()
+            chunk = prior_hidden_states.shape[1] // sp_world_size
+            prior_hidden_states = prior_hidden_states[
+                :, rank * chunk : (rank + 1) * chunk, :
+            ]
+        hidden_states = hidden_states + prior_hidden_states
+
+        temb = self.time_condition_embed(
+            timestep, target_size, crop_coords, hidden_states.dtype
+        )
+        temb = F.silu(temb)
+
+        # 3. Transformer blocks
+        for idx, block in enumerate(self.transformer_blocks):
+            hidden_states, encoder_hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                temb,
+                image_rotary_emb,
+                attention_mask,
+                attention_kwargs,
+                kv_cache=kv_caches[idx] if kv_caches is not None else None,
+            )
+
+        # 4. Output norm & projection
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 5. Unpatchify
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_height, post_patch_width, -1, p, p
+        )
+        output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
+
+        return output.float()
+        # float()
+        # reference: https://github.com/zRzRzRzRzRzRzR/diffusers/blob/6cfc83b4abc5b083fef56a18ec4700f48ba3aaba/src/diffusers/pipelines/glm_image/pipeline_glm_image.py#L737
+
+
+EntryClass = GlmImageTransformer2DModel

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

@@ -0,0 +1,974 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from Helios diffusers transformer:
+# https://github.com/BestWishYsh/Helios
+"""
+Helios Transformer 3D model for video generation.
+
+Implements the HeliosTransformer3DModel with multi-term memory patches,
+3D rotary position embeddings, and per-block scale-shift modulation.
+"""
+
+import math
+from functools import lru_cache
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_sp_world_size,
+    get_tp_world_size,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    LayerNorm,
+    LayerNormScaleShift,
+    RMSNorm,
+    tensor_parallel_rms_norm,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.visual_embedding import (
+    ModulateProjection,
+    PatchEmbed,
+    TimestepEmbedder,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Utility functions
+# ---------------------------------------------------------------------------
+
+
+def pad_for_3d_conv(x, kernel_size):
+    """Pad input to make it divisible by kernel_size using replicate mode."""
+    b, c, t, h, w = x.shape
+    pt, ph, pw = kernel_size
+    pad_t = (pt - (t % pt)) % pt
+    pad_h = (ph - (h % ph)) % ph
+    pad_w = (pw - (w % pw)) % pw
+    return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate")
+
+
+def center_down_sample_3d(x, kernel_size):
+    """Average pooling for 3D downsampling."""
+    return F.avg_pool3d(x, kernel_size, stride=kernel_size)
+
+
+def apply_rotary_emb_transposed(hidden_states, freqs_cis):
+    """Apply rotary positional embeddings with transposed cos/sin format."""
+    x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
+    cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
+    out = torch.empty_like(hidden_states)
+    out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2]
+    out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2]
+    return out.type_as(hidden_states)
+
+
+# ---------------------------------------------------------------------------
+# Output norm
+# ---------------------------------------------------------------------------
+
+
+class HeliosOutputNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+        self.norm = LayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
+        )
+
+    def forward(self, hidden_states, temb, original_context_length):
+        temb = temb[:, -original_context_length:, :]
+        shift, scale = (
+            self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)
+        ).chunk(2, dim=2)
+        shift = shift.squeeze(2).to(hidden_states.device)
+        scale = scale.squeeze(2).to(hidden_states.device)
+        hidden_states = hidden_states[:, -original_context_length:, :]
+        hidden_states = self.norm(hidden_states, shift, scale)
+        return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Rotary Positional Embedding (3D)
+# ---------------------------------------------------------------------------
+
+
+class HeliosRotaryPosEmbed(nn.Module):
+    """3D rotary position embeddings for (time, height, width)."""
+
+    def __init__(self, rope_dim, theta):
+        super().__init__()
+        self.DT, self.DY, self.DX = rope_dim
+        self.theta = theta
+        # Store as plain attributes (not buffers) to avoid meta-device issues
+        # during FSDP loading. They'll be re-created on the correct device in forward.
+        self._freqs_base_t = None
+        self._freqs_base_y = None
+        self._freqs_base_x = None
+
+    def _get_freqs_base(self, dim):
+        return 1.0 / (
+            self.theta
+            ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)
+        )
+
+    def _ensure_freqs_base(self, device):
+        """Lazily create frequency bases on the correct device."""
+        if self._freqs_base_t is None or self._freqs_base_t.device != device:
+            self._freqs_base_t = self._get_freqs_base(self.DT).to(device)
+            self._freqs_base_y = self._get_freqs_base(self.DY).to(device)
+            self._freqs_base_x = self._get_freqs_base(self.DX).to(device)
+
+    @torch.no_grad()
+    def get_frequency_batched(self, freqs_base, pos):
+        freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos)
+        freqs = freqs.repeat_interleave(2, dim=0)
+        return freqs.cos(), freqs.sin()
+
+    @torch.no_grad()
+    @lru_cache(maxsize=32)
+    def _get_spatial_meshgrid(self, height, width, device_str):
+        device = torch.device(device_str)
+        grid_y_coords = torch.arange(height, device=device, dtype=torch.float32)
+        grid_x_coords = torch.arange(width, device=device, dtype=torch.float32)
+        grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij")
+        return grid_y, grid_x
+
+    @torch.no_grad()
+    def forward(self, frame_indices, height, width, device):
+        self._ensure_freqs_base(device)
+        batch_size = frame_indices.shape[0]
+        num_frames = frame_indices.shape[1]
+
+        frame_indices = frame_indices.to(device=device, dtype=torch.float32)
+        grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device))
+
+        grid_t = frame_indices[:, :, None, None].expand(
+            batch_size, num_frames, height, width
+        )
+        grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1)
+        grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1)
+
+        freqs_cos_t, freqs_sin_t = self.get_frequency_batched(
+            self._freqs_base_t, grid_t
+        )
+        freqs_cos_y, freqs_sin_y = self.get_frequency_batched(
+            self._freqs_base_y, grid_y_batch
+        )
+        freqs_cos_x, freqs_sin_x = self.get_frequency_batched(
+            self._freqs_base_x, grid_x_batch
+        )
+
+        result = torch.cat(
+            [
+                freqs_cos_t,
+                freqs_cos_y,
+                freqs_cos_x,
+                freqs_sin_t,
+                freqs_sin_y,
+                freqs_sin_x,
+            ],
+            dim=0,
+        )
+        return result.permute(1, 0, 2, 3, 4)
+
+
+# ---------------------------------------------------------------------------
+# Condition Embedder
+# ---------------------------------------------------------------------------
+
+
+class HeliosTimeTextEmbedding(nn.Module):
+    """Condition embedder combining timestep and text embeddings."""
+
+    def __init__(self, dim, time_freq_dim, time_proj_dim, text_embed_dim):
+        super().__init__()
+        self.time_embedder = TimestepEmbedder(
+            dim, frequency_embedding_size=time_freq_dim, act_layer="silu"
+        )
+        self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu")
+        self.text_embedder = MLP(
+            text_embed_dim, dim, dim, bias=True, act_type="gelu_pytorch_tanh"
+        )
+
+    def forward(
+        self, timestep, encoder_hidden_states, is_return_encoder_hidden_states=True
+    ):
+        temb = self.time_embedder(timestep)
+        timestep_proj = self.time_modulation(temb)
+
+        if encoder_hidden_states is not None and is_return_encoder_hidden_states:
+            encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+
+        return temb, timestep_proj, encoder_hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Self-Attention for Helios
+# ---------------------------------------------------------------------------
+
+
+class HeliosSelfAttention(nn.Module):
+    """Self-attention with RMSNorm Q/K, optional history key amplification."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        is_amplify_history: bool = False,
+        history_scale_mode: str = "per_head",
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        tp_size = get_tp_world_size()
+        self.local_num_heads = divide(num_heads, tp_size)
+
+        self.to_q = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_out = RowParallelLinear(
+            dim, dim, bias=True, reduce_results=True, quant_config=quant_config
+        )
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+        self.tp_rmsnorm = tp_size > 1
+
+        self.attn = USPAttention(
+            num_heads=self.local_num_heads,
+            head_size=self.head_dim,
+            causal=False,
+            is_cross_attention=False,
+        )
+
+        self.is_amplify_history = is_amplify_history
+        if is_amplify_history:
+            if history_scale_mode == "scalar":
+                self.history_key_scale = nn.Parameter(torch.ones(1))
+            elif history_scale_mode == "per_head":
+                self.history_key_scale = nn.Parameter(torch.ones(num_heads))
+            else:
+                raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}")
+            self.history_scale_mode = history_scale_mode
+            self.max_scale = 10.0
+
+    def forward(self, hidden_states, rotary_emb=None, original_context_length=None):
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(hidden_states)
+        v, _ = self.to_v(hidden_states)
+
+        if self.tp_rmsnorm:
+            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 = q.unflatten(2, (self.local_num_heads, self.head_dim))
+        k = k.unflatten(2, (self.local_num_heads, self.head_dim))
+        v = v.unflatten(2, (self.local_num_heads, self.head_dim))
+
+        if rotary_emb is not None:
+            q = apply_rotary_emb_transposed(q, rotary_emb)
+            k = apply_rotary_emb_transposed(k, rotary_emb)
+
+        history_seq_len = (
+            hidden_states.shape[1] - original_context_length
+            if original_context_length is not None
+            else 0
+        )
+
+        if self.is_amplify_history and original_context_length is not None:
+            if history_seq_len > 0:
+                scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * (
+                    self.max_scale - 1.0
+                )
+                if self.history_scale_mode == "per_head":
+                    scale_key = scale_key.view(1, 1, -1, 1)
+                k = torch.cat(
+                    [k[:, :history_seq_len] * scale_key, k[:, history_seq_len:]],
+                    dim=1,
+                )
+
+        x = self.attn(q, k, v, num_replicated_prefix=history_seq_len)
+        x = x.flatten(2)
+        x, _ = self.to_out(x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Cross-Attention for Helios
+# ---------------------------------------------------------------------------
+
+
+class HeliosCrossAttention(nn.Module):
+    """Cross-attention with RMSNorm Q/K normalization."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        tp_size = get_tp_world_size()
+        self.local_num_heads = divide(num_heads, tp_size)
+
+        self.to_q = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.to_out = RowParallelLinear(
+            dim, dim, bias=True, reduce_results=True, quant_config=quant_config
+        )
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+        self.tp_rmsnorm = tp_size > 1
+
+        self.attn = USPAttention(
+            num_heads=self.local_num_heads,
+            head_size=self.head_dim,
+            causal=False,
+            skip_sequence_parallel=True,
+        )
+
+    def project_kv(self, encoder_hidden_states):
+        """Project encoder states to this block's cross-attn (key, value)."""
+        k, _ = self.to_k(encoder_hidden_states)
+        v, _ = self.to_v(encoder_hidden_states)
+        if self.tp_rmsnorm:
+            k = tensor_parallel_rms_norm(k, self.norm_k)
+        else:
+            k = self.norm_k(k)
+        k = k.unflatten(2, (self.local_num_heads, self.head_dim))
+        v = v.unflatten(2, (self.local_num_heads, self.head_dim))
+        return k, v
+
+    def forward(
+        self, hidden_states, encoder_hidden_states=None, encoder_key_value=None
+    ):
+        q, _ = self.to_q(hidden_states)
+        if self.tp_rmsnorm:
+            q = tensor_parallel_rms_norm(q, self.norm_q)
+        else:
+            q = self.norm_q(q)
+        q = q.unflatten(2, (self.local_num_heads, self.head_dim))
+
+        if encoder_key_value is None:
+            if encoder_hidden_states is None:
+                raise ValueError(
+                    "encoder_hidden_states is required when encoder_key_value"
+                    " is not provided."
+                )
+            encoder_key_value = self.project_kv(encoder_hidden_states)
+        k, v = encoder_key_value
+
+        x = self.attn(q, k, v)
+        x = x.flatten(2)
+        x, _ = self.to_out(x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Transformer Block
+# ---------------------------------------------------------------------------
+
+
+class HeliosTransformerBlock(nn.Module):
+    """
+    Single transformer block with self-attention, cross-attention, FFN,
+    and scale-shift modulation from timestep embeddings.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        cross_attn_norm: bool = True,
+        eps: float = 1e-6,
+        guidance_cross_attn: bool = True,
+        is_amplify_history: bool = False,
+        history_scale_mode: str = "per_head",
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__()
+
+        # 1. Self-attention
+        self.norm1 = LayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
+        )
+        self.attn1 = HeliosSelfAttention(
+            dim=dim,
+            num_heads=num_heads,
+            eps=eps,
+            is_amplify_history=is_amplify_history,
+            history_scale_mode=history_scale_mode,
+            quant_config=quant_config,
+        )
+
+        # 2. Cross-attention
+        self.attn2 = HeliosCrossAttention(
+            dim=dim,
+            num_heads=num_heads,
+            eps=eps,
+            quant_config=quant_config,
+        )
+        self.self_attn_residual_norm = (
+            LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=torch.float32)
+            if cross_attn_norm
+            else nn.Identity()
+        )
+
+        # 3. Feed-forward
+        self.ffn = MLP(
+            dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config
+        )
+        self.norm3 = LayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
+        )
+
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+        # 4. Guidance cross-attention flag
+        self.guidance_cross_attn = guidance_cross_attn
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states,
+        temb,
+        rotary_emb,
+        original_context_length=None,
+        cross_attn_key_value=None,
+    ):
+        if temb.ndim == 4:
+            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
+                self.scale_shift_table.unsqueeze(0) + temb.float()
+            ).chunk(6, dim=2)
+            shift_msa = shift_msa.squeeze(2)
+            scale_msa = scale_msa.squeeze(2)
+            gate_msa = gate_msa.squeeze(2)
+            c_shift_msa = c_shift_msa.squeeze(2)
+            c_scale_msa = c_scale_msa.squeeze(2)
+            c_gate_msa = c_gate_msa.squeeze(2)
+        else:
+            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
+                self.scale_shift_table + temb.float()
+            ).chunk(6, dim=1)
+
+        # 1. Self-attention
+        norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa)
+        attn_output = self.attn1(
+            norm_hidden_states, rotary_emb, original_context_length
+        )
+        hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(
+            hidden_states
+        )
+
+        # 2. Cross-attention
+        if self.guidance_cross_attn:
+            history_seq_len = hidden_states.shape[1] - original_context_length
+            history_hidden_states, current_hidden_states = torch.split(
+                hidden_states, [history_seq_len, original_context_length], dim=1
+            )
+            norm_hidden_states = self.self_attn_residual_norm(
+                current_hidden_states.float()
+            ).type_as(current_hidden_states)
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states,
+                encoder_key_value=cross_attn_key_value,
+            )
+            current_hidden_states = current_hidden_states + attn_output
+            hidden_states = torch.cat(
+                [history_hidden_states, current_hidden_states], dim=1
+            )
+        else:
+            norm_hidden_states = self.self_attn_residual_norm(
+                hidden_states.float()
+            ).type_as(hidden_states)
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states,
+                encoder_key_value=cross_attn_key_value,
+            )
+            hidden_states = hidden_states + attn_output
+
+        # 3. Feed-forward
+        norm_hidden_states = self.norm3(hidden_states, c_shift_msa, c_scale_msa)
+        ff_output = self.ffn(norm_hidden_states)
+        hidden_states = (
+            hidden_states.float() + ff_output.float() * c_gate_msa
+        ).type_as(hidden_states)
+
+        return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Main model
+# ---------------------------------------------------------------------------
+
+
+class HeliosTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """
+    Helios Transformer 3D model for video generation.
+
+    Implements multi-scale history patches, 3D RoPE, and chunked denoising
+    with zero_history_timestep and guidance_cross_attn.
+    """
+
+    _fsdp_shard_conditions = HeliosConfig()._fsdp_shard_conditions
+    _compile_conditions = HeliosConfig()._compile_conditions
+    _supported_attention_backends = HeliosConfig()._supported_attention_backends
+    param_names_mapping = HeliosConfig().param_names_mapping
+    reverse_param_names_mapping = HeliosConfig().reverse_param_names_mapping
+    lora_param_names_mapping = HeliosConfig().lora_param_names_mapping
+
+    def __init__(
+        self,
+        config: HeliosConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_channels_latents = config.num_channels_latents
+        self.patch_size = config.patch_size
+        self.text_len = config.text_len
+        self.inner_dim = inner_dim
+
+        # Helios-specific config
+        self.zero_history_timestep = config.zero_history_timestep
+        self.has_multi_term_memory_patch = config.has_multi_term_memory_patch
+        self.guidance_cross_attn = config.guidance_cross_attn
+
+        # 1. Patch & position embedding
+        self.patch_embedding = PatchEmbed(
+            in_chans=config.in_channels,
+            embed_dim=inner_dim,
+            patch_size=config.patch_size,
+            flatten=False,
+        )
+
+        # 2. Rotary position embeddings
+        self.rope = HeliosRotaryPosEmbed(
+            rope_dim=config.rope_dim, theta=config.rope_theta
+        )
+
+        # 3. Multi-term memory patches
+        if self.has_multi_term_memory_patch:
+            self.patch_short = nn.Conv3d(
+                config.in_channels,
+                inner_dim,
+                kernel_size=config.patch_size,
+                stride=config.patch_size,
+            )
+            self.patch_mid = nn.Conv3d(
+                config.in_channels,
+                inner_dim,
+                kernel_size=tuple(2 * p for p in config.patch_size),
+                stride=tuple(2 * p for p in config.patch_size),
+            )
+            self.patch_long = nn.Conv3d(
+                config.in_channels,
+                inner_dim,
+                kernel_size=tuple(4 * p for p in config.patch_size),
+                stride=tuple(4 * p for p in config.patch_size),
+            )
+
+        # 4. Condition embeddings
+        self.condition_embedder = HeliosTimeTextEmbedding(
+            dim=inner_dim,
+            time_freq_dim=config.freq_dim,
+            time_proj_dim=inner_dim * 6,
+            text_embed_dim=config.text_dim,
+        )
+
+        # 5. Transformer blocks
+        self.blocks = nn.ModuleList(
+            [
+                HeliosTransformerBlock(
+                    dim=inner_dim,
+                    ffn_dim=config.ffn_dim,
+                    num_heads=config.num_attention_heads,
+                    cross_attn_norm=config.cross_attn_norm,
+                    eps=config.eps,
+                    guidance_cross_attn=config.guidance_cross_attn,
+                    is_amplify_history=config.is_amplify_history,
+                    history_scale_mode=config.history_scale_mode,
+                    quant_config=quant_config,
+                )
+                for _ in range(config.num_layers)
+            ]
+        )
+
+        # 6. Output norm & projection
+        self.norm_out = HeliosOutputNorm(inner_dim, config.eps)
+        self.proj_out = ColumnParallelLinear(
+            inner_dim,
+            config.out_channels * math.prod(config.patch_size),
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
+        )
+
+        self.cnt = 0
+        self.__post_init__()
+        self.layer_names = ["blocks"]
+        self.sp_size = get_sp_world_size()
+
+    # Cross-attention K/V cache.
+    #
+    # Text conditioning is constant across the denoise loop, so the text
+    # projection and every block's cross-attn K/V are computed once per request
+    # (keyed by encoder-tensor identity) and reused across steps.
+
+    @staticmethod
+    def _request_cache(forward_batch, name):
+        """Per-request cache dict on ``forward_batch.extra``.
+
+        Returns None (-> caller recomputes, caching disabled) when there is no
+        forward batch or gradients are enabled."""
+        if forward_batch is None or torch.is_grad_enabled():
+            return None
+        extra = getattr(forward_batch, "extra", None)
+        return None if extra is None else extra.setdefault(name, {})
+
+    @staticmethod
+    def _tensor_key(t):
+        """Identity key for ``t``; equal only for the same underlying tensor."""
+        return (
+            t.data_ptr(),
+            tuple(t.shape),
+            tuple(t.stride()),
+            t.dtype,
+            t.device.type,
+            t.device.index,
+        )
+
+    def _get_cross_attn_key_values(self, encoder_hidden_states, forward_batch):
+        """Per-block cross-attn (key, value) for ``encoder_hidden_states``.
+
+        Cached per request, keyed on the encoder tensor's identity
+        (``_tensor_key``). The same object — ``batch.prompt_embeds`` — is passed
+        every denoise step, so the key is stable and steps after the first hit
+        the cache.
+        """
+        cache = self._request_cache(forward_batch, "helios_cross_attn_kv")
+        key = self._tensor_key(encoder_hidden_states) if cache is not None else None
+        kvs = cache.get(key) if key is not None else None
+        if kvs is None:
+            projected = self.condition_embedder.text_embedder(encoder_hidden_states)
+            kvs = [block.attn2.project_kv(projected) for block in self.blocks]
+            if key is not None:
+                cache[key] = kvs
+        return kvs
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        # Stage 1 history inputs
+        indices_hidden_states=None,
+        indices_latents_history_short=None,
+        indices_latents_history_mid=None,
+        indices_latents_history_long=None,
+        latents_history_short=None,
+        latents_history_mid=None,
+        latents_history_long=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        # Check if sequence parallelism is enabled
+        forward_batch = get_forward_context().forward_batch
+        if forward_batch is not None:
+            sequence_shard_enabled = (
+                forward_batch.enable_sequence_shard and self.sp_size > 1
+            )
+        else:
+            sequence_shard_enabled = False
+
+        batch_size = hidden_states.shape[0]
+        p_t, p_h, p_w = self.patch_size
+
+        # 1. Patch embed the noisy latents
+        hidden_states = self.patch_embedding(hidden_states)
+        (
+            _,
+            _,
+            post_patch_num_frames,
+            post_patch_height,
+            post_patch_width,
+        ) = hidden_states.shape
+
+        if indices_hidden_states is None:
+            indices_hidden_states = (
+                torch.arange(0, post_patch_num_frames)
+                .unsqueeze(0)
+                .expand(batch_size, -1)
+            )
+
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        # 2. Compute rotary embeddings
+        rotary_emb = self.rope(
+            frame_indices=indices_hidden_states,
+            height=post_patch_height,
+            width=post_patch_width,
+            device=hidden_states.device,
+        )
+        rotary_emb = rotary_emb.flatten(2).transpose(1, 2)
+        original_context_length = hidden_states.shape[1]
+
+        # Sequence parallelism: shard current tokens and RoPE across SP ranks
+        seq_shard_pad = 0
+        if sequence_shard_enabled:
+            sp_rank = get_sp_group().rank_in_group
+            seq_len = hidden_states.shape[1]
+            if seq_len % self.sp_size != 0:
+                seq_shard_pad = self.sp_size - (seq_len % self.sp_size)
+                hs_pad = torch.zeros(
+                    batch_size,
+                    seq_shard_pad,
+                    hidden_states.shape[2],
+                    dtype=hidden_states.dtype,
+                    device=hidden_states.device,
+                )
+                re_pad = torch.zeros(
+                    batch_size,
+                    seq_shard_pad,
+                    rotary_emb.shape[2],
+                    dtype=rotary_emb.dtype,
+                    device=rotary_emb.device,
+                )
+                hidden_states = torch.cat([hidden_states, hs_pad], dim=1)
+                rotary_emb = torch.cat([rotary_emb, re_pad], dim=1)
+            local_seq_len = hidden_states.shape[1] // self.sp_size
+            hidden_states = hidden_states.view(
+                batch_size, self.sp_size, local_seq_len, -1
+            )[:, sp_rank, :, :].contiguous()
+            rotary_emb = rotary_emb.view(batch_size, self.sp_size, local_seq_len, -1)[
+                :, sp_rank, :, :
+            ].contiguous()
+            effective_context_length = local_seq_len
+        else:
+            effective_context_length = original_context_length
+
+        # 3. Process short history
+        if (
+            latents_history_short is not None
+            and indices_latents_history_short is not None
+        ):
+            latents_history_short = latents_history_short.to(hidden_states)
+            latents_history_short = self.patch_short(latents_history_short)
+            _, _, _, H1, W1 = latents_history_short.shape
+            latents_history_short = latents_history_short.flatten(2).transpose(1, 2)
+
+            rotary_emb_history_short = self.rope(
+                frame_indices=indices_latents_history_short,
+                height=H1,
+                width=W1,
+                device=latents_history_short.device,
+            )
+            rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(
+                1, 2
+            )
+            hidden_states = torch.cat([latents_history_short, hidden_states], dim=1)
+            rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1)
+
+        # 4. Process mid history
+        if latents_history_mid is not None and indices_latents_history_mid is not None:
+            latents_history_mid = latents_history_mid.to(hidden_states)
+            latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4))
+            latents_history_mid = self.patch_mid(latents_history_mid)
+            latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2)
+
+            rotary_emb_history_mid = self.rope(
+                frame_indices=indices_latents_history_mid,
+                height=H1,
+                width=W1,
+                device=latents_history_mid.device,
+            )
+            rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2))
+            rotary_emb_history_mid = center_down_sample_3d(
+                rotary_emb_history_mid, (2, 2, 2)
+            )
+            rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2)
+
+            hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1)
+            rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1)
+
+        # 5. Process long history
+        if (
+            latents_history_long is not None
+            and indices_latents_history_long is not None
+        ):
+            latents_history_long = latents_history_long.to(hidden_states)
+            latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8))
+            latents_history_long = self.patch_long(latents_history_long)
+            latents_history_long = latents_history_long.flatten(2).transpose(1, 2)
+
+            rotary_emb_history_long = self.rope(
+                frame_indices=indices_latents_history_long,
+                height=H1,
+                width=W1,
+                device=latents_history_long.device,
+            )
+            rotary_emb_history_long = pad_for_3d_conv(
+                rotary_emb_history_long, (4, 4, 4)
+            )
+            rotary_emb_history_long = center_down_sample_3d(
+                rotary_emb_history_long, (4, 4, 4)
+            )
+            rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2)
+
+            hidden_states = torch.cat([latents_history_long, hidden_states], dim=1)
+            rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1)
+
+        history_context_length = hidden_states.shape[1] - effective_context_length
+
+        # 6. Compute condition embeddings
+        if indices_hidden_states is not None and self.zero_history_timestep:
+            timestep_t0 = torch.zeros(
+                (1,), dtype=timestep.dtype, device=timestep.device
+            )
+            temb_t0, timestep_proj_t0, _ = self.condition_embedder(
+                timestep_t0,
+                encoder_hidden_states,
+                is_return_encoder_hidden_states=False,
+            )
+            temb_t0 = temb_t0.unsqueeze(1).expand(
+                batch_size, history_context_length, -1
+            )
+            timestep_proj_t0 = (
+                timestep_proj_t0.unflatten(-1, (6, -1))
+                .view(1, 6, 1, -1)
+                .expand(batch_size, -1, history_context_length, -1)
+            )
+
+        # Take only the time embeddings (temb, timestep_proj); skip the text
+        # projection (is_return_encoder_hidden_states=False) since it is computed
+        # once per request and cached by _get_cross_attn_key_values below.
+        temb, timestep_proj, _ = self.condition_embedder(
+            timestep, encoder_hidden_states, is_return_encoder_hidden_states=False
+        )
+        cross_attn_key_values = self._get_cross_attn_key_values(
+            encoder_hidden_states, forward_batch
+        )
+        timestep_proj = timestep_proj.unflatten(-1, (6, -1))
+
+        if indices_hidden_states is not None and not self.zero_history_timestep:
+            main_repeat_size = hidden_states.shape[1]
+        else:
+            main_repeat_size = effective_context_length
+        temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1)
+        timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand(
+            batch_size, 6, main_repeat_size, -1
+        )
+
+        if indices_hidden_states is not None and self.zero_history_timestep:
+            temb = torch.cat([temb_t0, temb], dim=1)
+            timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2)
+
+        if timestep_proj.ndim == 4:
+            timestep_proj = timestep_proj.permute(0, 2, 1, 3)
+
+        # 7. Transformer blocks
+        hidden_states = hidden_states.contiguous()
+        encoder_hidden_states = encoder_hidden_states.contiguous()
+        rotary_emb = rotary_emb.contiguous()
+
+        for block, key_value in zip(self.blocks, cross_attn_key_values):
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                timestep_proj,
+                rotary_emb,
+                effective_context_length,
+                cross_attn_key_value=key_value,
+            )
+
+        self.cnt += 1
+
+        # SP: all-gather current tokens before output
+        if sequence_shard_enabled:
+            current_tokens = hidden_states[:, -local_seq_len:, :].contiguous()
+            current_tokens = sequence_model_parallel_all_gather(current_tokens, dim=1)
+            if seq_shard_pad > 0:
+                current_tokens = current_tokens[:, :original_context_length, :]
+            hidden_states = current_tokens
+            # Re-create temb for norm_out (all current tokens share same timestep)
+            temb = temb[:, :1, :].expand(batch_size, original_context_length, -1)
+
+        # 8. Output norm & projection
+        hidden_states = self.norm_out(hidden_states, temb, original_context_length)
+        hidden_states, _ = self.proj_out(hidden_states)
+
+        # 9. Unpatchify
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            post_patch_num_frames,
+            post_patch_height,
+            post_patch_width,
+            p_t,
+            p_h,
+            p_w,
+            -1,
+        )
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return output
+
+
+EntryClass = HeliosTransformer3DModel

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

@@ -0,0 +1,575 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_tp_world_size,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.layers.attention import (
+    USPAttention,
+    build_varlen_mask_meta,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8 import (
+    WeightOnlyFP8ColumnParallelLinear,
+    WeightOnlyFP8Linear,
+    WeightOnlyFP8MergedColumnParallelLinear,
+    WeightOnlyFP8RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
+    Qwen3VLTextRotaryEmbedding,
+    qwen3_apply_rotary_pos_emb,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT
+
+OUTPUT_IMAGE_INDICATOR = 2
+LLM_TOKEN_INDICATOR = 3
+
+
+class Ideogram4RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x, self.weight.shape, self.weight, self.eps)
+
+
+class Ideogram4QuantizedLinear(ReplicatedLinear):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x)[0]
+
+
+class Ideogram4ColumnParallelLinear(ColumnParallelLinear):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x)[0]
+
+
+class Ideogram4MergedColumnParallelLinear(MergedColumnParallelLinear):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x)[0]
+
+
+class Ideogram4RowParallelLinear(RowParallelLinear):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x)[0]
+
+
+def _tp_size() -> int:
+    return get_tp_world_size() if model_parallel_is_initialized() else 1
+
+
+def _linear(
+    in_features: int,
+    out_features: int,
+    bias: bool = True,
+    quant_config: QuantizationConfig | None = None,
+    prefix: str = "",
+    gather_output: bool = True,
+):
+    tp_size = _tp_size()
+    use_column_parallel = tp_size > 1 and out_features % tp_size == 0
+    if quant_config is None:
+        if use_column_parallel:
+            return WeightOnlyFP8ColumnParallelLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                gather_output=gather_output,
+            )
+        return WeightOnlyFP8Linear(in_features, out_features, bias=bias)
+    if use_column_parallel:
+        return Ideogram4ColumnParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+            gather_output=gather_output,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+    return Ideogram4QuantizedLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+    )
+
+
+def _merged_column_linear(
+    in_features: int,
+    output_sizes: list[int],
+    bias: bool = True,
+    quant_config: QuantizationConfig | None = None,
+    prefix: str = "",
+):
+    tp_size = _tp_size()
+    use_column_parallel = tp_size > 1 and all(
+        output_size % tp_size == 0 for output_size in output_sizes
+    )
+    out_features = sum(output_sizes)
+    if quant_config is None:
+        if use_column_parallel:
+            return WeightOnlyFP8MergedColumnParallelLinear(
+                in_features,
+                output_sizes,
+                bias=bias,
+                gather_output=False,
+            )
+        return WeightOnlyFP8Linear(in_features, out_features, bias=bias)
+    if use_column_parallel:
+        return Ideogram4MergedColumnParallelLinear(
+            in_features,
+            output_sizes,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+    return Ideogram4QuantizedLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+    )
+
+
+def _row_linear(
+    in_features: int,
+    out_features: int,
+    bias: bool = True,
+    quant_config: QuantizationConfig | None = None,
+    prefix: str = "",
+):
+    tp_size = _tp_size()
+    use_row_parallel = tp_size > 1 and in_features % tp_size == 0
+    if quant_config is None:
+        if use_row_parallel:
+            return WeightOnlyFP8RowParallelLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                input_is_parallel=True,
+            )
+        return WeightOnlyFP8Linear(in_features, out_features, bias=bias)
+    if use_row_parallel:
+        return Ideogram4RowParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+    return Ideogram4QuantizedLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+    )
+
+
+class Ideogram4Attention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        eps: float,
+        supported_attention_backends,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        tp_size = _tp_size()
+        assert num_heads % tp_size == 0
+        self.local_num_heads = divide(num_heads, tp_size)
+        self.qkv = _merged_column_linear(
+            hidden_size,
+            [hidden_size, hidden_size, hidden_size],
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv",
+        )
+        self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps)
+        self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps)
+        self.attn = USPAttention(
+            num_heads=self.local_num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+        self.o = _row_linear(
+            hidden_size,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o",
+        )
+
+    def forward(self, x, cos, sin, attn_mask, attn_mask_meta):
+        batch_size, seq_len, _ = x.shape
+        qkv = self.qkv(x).view(
+            batch_size, seq_len, 3, self.local_num_heads, self.head_dim
+        )
+        q, k, v = qkv.unbind(dim=2)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
+        q, k = qwen3_apply_rotary_pos_emb(q, k, cos, sin)
+        out = self.attn(q, k, v, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta)
+        out = out.reshape(batch_size, seq_len, self.local_num_heads * self.head_dim)
+        return self.o(out)
+
+
+class Ideogram4MLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.w1 = _linear(
+            dim,
+            hidden_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.w1",
+            gather_output=False,
+        )
+        self.w2 = _row_linear(
+            hidden_dim,
+            dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.w2",
+        )
+        self.w3 = _linear(
+            dim,
+            hidden_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.w3",
+            gather_output=False,
+        )
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class Ideogram4TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        intermediate_size,
+        num_heads,
+        norm_eps,
+        adaln_dim,
+        supported_attention_backends,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.attention = Ideogram4Attention(
+            hidden_size,
+            num_heads,
+            eps=1e-5,
+            supported_attention_backends=supported_attention_backends,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attention",
+        )
+        self.feed_forward = Ideogram4MLP(
+            hidden_size,
+            intermediate_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.feed_forward",
+        )
+        self.attention_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+        self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+        self.attention_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+        self.ffn_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+        self.adaln_modulation = _linear(
+            adaln_dim,
+            4 * hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.adaln_modulation",
+        )
+
+    def forward(self, x, cos, sin, adaln_input, attn_mask, attn_mask_meta):
+        scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaln_modulation(
+            adaln_input
+        ).chunk(4, dim=-1)
+        gate_msa = torch.tanh(gate_msa)
+        gate_mlp = torch.tanh(gate_mlp)
+        attn_out = self.attention(
+            self.attention_norm1(x) * (1.0 + scale_msa),
+            cos=cos,
+            sin=sin,
+            attn_mask=attn_mask,
+            attn_mask_meta=attn_mask_meta,
+        )
+        x = x + gate_msa * self.attention_norm2(attn_out)
+        x = x + gate_mlp * self.ffn_norm2(
+            self.feed_forward(self.ffn_norm1(x) * (1.0 + scale_mlp))
+        )
+        return x
+
+
+def _sinusoidal_embedding(t: torch.Tensor, dim: int, scale: float = 1e4):
+    t = t.to(torch.float32)
+    half = dim // 2
+    freq = math.log(scale) / (half - 1)
+    freq = torch.exp(torch.arange(half, dtype=torch.float32, device=t.device) * -freq)
+    emb = t.unsqueeze(-1) * freq
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+    if dim % 2 == 1:
+        emb = F.pad(emb, (0, 1))
+    return emb
+
+
+class Ideogram4EmbedScalar(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        input_range: tuple[float, float],
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.range_min, self.range_max = input_range
+        self.mlp_in = _linear(
+            dim,
+            dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp_in",
+        )
+        self.mlp_out = _linear(
+            dim,
+            dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp_out",
+        )
+
+    def forward(self, x):
+        compute_dtype = x.dtype
+        x = x.to(torch.float32)
+        scaled = 1e4 * (x - self.range_min) / (self.range_max - self.range_min)
+        emb = _sinusoidal_embedding(scaled, self.dim).to(compute_dtype)
+        return self.mlp_out(F.silu(self.mlp_in(emb)))
+
+
+class Ideogram4FinalLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        out_channels: int,
+        adaln_dim: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+        self.linear = _linear(
+            hidden_size,
+            out_channels,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
+        self.adaln_modulation = _linear(
+            adaln_dim,
+            hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.adaln_modulation",
+        )
+
+    def forward(self, x, c):
+        scale = 1.0 + self.adaln_modulation(F.silu(c))
+        return self.linear(self.norm_final(x) * scale)
+
+
+class Ideogram4Transformer2DModel(BaseDiT):
+    _repeated_blocks = ["Ideogram4TransformerBlock"]
+    _fsdp_shard_conditions = Ideogram4DiTConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = Ideogram4DiTConfig().arch_config._compile_conditions
+    _supported_attention_backends = (
+        Ideogram4DiTConfig().arch_config._supported_attention_backends
+    )
+    param_names_mapping = {}
+    reverse_param_names_mapping = {}
+
+    def __init__(
+        self,
+        config: Ideogram4DiTConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(config, hf_config, **kwargs)
+        cfg = config.arch_config
+        self._supported_attention_backends = cfg._supported_attention_backends
+        hidden_size = cfg.num_attention_heads * cfg.attention_head_dim
+        self.hidden_size = hidden_size
+        self.num_attention_heads = cfg.num_attention_heads
+        self.num_channels_latents = cfg.in_channels
+        self.input_proj = _linear(
+            cfg.in_channels,
+            hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix="input_proj",
+        )
+        self.llm_cond_norm = Ideogram4RMSNorm(cfg.llm_features_dim, eps=1e-6)
+        self.llm_cond_proj = _linear(
+            cfg.llm_features_dim,
+            hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix="llm_cond_proj",
+        )
+        self.t_embedding = Ideogram4EmbedScalar(
+            hidden_size,
+            input_range=(0.0, 1.0),
+            quant_config=quant_config,
+            prefix="t_embedding",
+        )
+        self.adaln_proj = _linear(
+            hidden_size,
+            cfg.adaln_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix="adaln_proj",
+        )
+        self.embed_image_indicator = nn.Embedding(2, hidden_size)
+        self.rotary_emb = Qwen3VLTextRotaryEmbedding(
+            head_dim=cfg.attention_head_dim,
+            rope_theta=cfg.rope_theta,
+            mrope_section=cfg.mrope_section,
+        )
+        self.layers = nn.ModuleList(
+            [
+                Ideogram4TransformerBlock(
+                    hidden_size=hidden_size,
+                    intermediate_size=cfg.intermediate_size,
+                    num_heads=cfg.num_attention_heads,
+                    norm_eps=cfg.norm_eps,
+                    adaln_dim=cfg.adaln_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                    quant_config=quant_config,
+                    prefix=f"layers.{i}",
+                )
+                for i in range(cfg.num_layers)
+            ]
+        )
+        self.final_layer = Ideogram4FinalLayer(
+            hidden_size=hidden_size,
+            out_channels=cfg.in_channels,
+            adaln_dim=cfg.adaln_dim,
+            quant_config=quant_config,
+            prefix="final_layer",
+        )
+
+    def post_load_weights(self) -> None:
+        if not self.rotary_emb.inv_freq.is_meta:
+            return
+        cfg = self.config.arch_config
+        inv_freq = 1.0 / (
+            cfg.rope_theta
+            ** (
+                torch.arange(
+                    0,
+                    cfg.attention_head_dim,
+                    2,
+                    dtype=torch.float32,
+                    device=self.input_proj.weight.device,
+                )
+                / cfg.attention_head_dim
+            )
+        )
+        self.rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(
+        self,
+        *,
+        llm_features: torch.Tensor,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        position_ids: torch.Tensor,
+        segment_ids: torch.Tensor,
+        indicator: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        attn_mask_meta: dict | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        param_dtype = self.embed_image_indicator.weight.dtype
+        x = x.to(param_dtype)
+        t = t.to(param_dtype)
+        llm_features = llm_features.to(param_dtype)
+        indicator = indicator.to(torch.long)
+        llm_token_mask = (indicator == LLM_TOKEN_INDICATOR).to(x.dtype).unsqueeze(-1)
+        output_image_mask = (
+            (indicator == OUTPUT_IMAGE_INDICATOR).to(x.dtype).unsqueeze(-1)
+        )
+        llm_features = llm_features * llm_token_mask
+        x = x * output_image_mask
+        x = self.input_proj(x) * output_image_mask
+        t_cond = self.t_embedding(t)
+        if t.dim() == 1:
+            t_cond = t_cond.unsqueeze(1)
+        adaln_input = F.silu(self.adaln_proj(t_cond))
+        llm_features = self.llm_cond_proj(self.llm_cond_norm(llm_features))
+        llm_features = llm_features * llm_token_mask
+        h = x + llm_features
+        h = h + self.embed_image_indicator(
+            (indicator == OUTPUT_IMAGE_INDICATOR).to(torch.long)
+        )
+        cos, sin = self.rotary_emb(h, position_ids)
+        cos = cos.unsqueeze(2)
+        sin = sin.unsqueeze(2)
+        # ideogram uses -1 padding; varlen meta enables fa packed attention
+        if attn_mask is None:
+            attn_mask = segment_ids > 0
+        if attn_mask_meta is None:
+            attn_mask_meta = build_varlen_mask_meta(attn_mask)
+        for layer in self.layers:
+            h = layer(
+                h,
+                cos=cos,
+                sin=sin,
+                adaln_input=adaln_input,
+                attn_mask=attn_mask,
+                attn_mask_meta=attn_mask_meta,
+            )
+        return self.final_layer(h, c=adaln_input).to(torch.float32)
+
+
+EntryClass = Ideogram4Transformer2DModel

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

@@ -0,0 +1,596 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from functools import lru_cache
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from sglang.multimodal_gen.configs.models.dits.joy_image import JoyImageDiTConfig
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_sp_group,
+    get_sp_world_size,
+    get_tp_world_size,
+    sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    LayerNormScaleShift,
+    RMSNorm,
+    apply_qk_norm_with_optional_rope,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import NDRotaryEmbedding
+from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.models.dits.wanvideo import WanTimeTextImageEmbedding
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+logger = init_logger(__name__)
+_MODULATION_FACTOR = 6
+
+
+def fused_add_gate(
+    residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor
+) -> torch.Tensor:
+    """Fused residual addition with gate.
+
+    Computes: residual + x * gate.unsqueeze(1)
+
+    This fuses the gate multiplication and residual addition to reduce
+    intermediate tensor allocations and memory bandwidth.
+
+    Args:
+        residual (torch.Tensor): The residual tensor to add to. Shape: (B, L, D)
+        x (torch.Tensor): The input tensor to be gated. Shape: (B, L, D)
+        gate (torch.Tensor): The gate tensor. Shape: (B, D)
+
+    Returns:
+        torch.Tensor: residual + x * gate.unsqueeze(1)
+    """
+    return torch.addcmul(residual, x, gate.unsqueeze(1))
+
+
+class ModulateWan(nn.Module):
+    """Modulation layer for WanX."""
+
+    def __init__(self, hidden_size: int, factor: int, dtype=None, device=None):
+        super().__init__()
+        self.factor = factor
+        self.modulate_table = nn.Parameter(
+            torch.zeros(1, factor, hidden_size, dtype=dtype, device=device)
+            / hidden_size**0.5,
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.modulate_table,
+            {
+                "input_dim": 1,
+                "output_dim": 2,
+            },
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if len(x.shape) != 3:
+            x = x.unsqueeze(1)
+        return [
+            o.squeeze(1) for o in (self.modulate_table + x).chunk(self.factor, dim=1)
+        ]
+
+
+class MMDoubleStreamBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        heads_num: int,
+        mlp_width_ratio: float,
+        mlp_act_type: str = "gelu_pytorch_tanh",
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.heads_num = heads_num
+        self.hidden_size = hidden_size
+        self.tp_size = get_tp_world_size()
+        self.local_heads_num = divide(self.heads_num, self.tp_size)
+        self.head_dim = self.hidden_size // self.heads_num
+        self.mlp_hidden_dim = int(self.hidden_size * mlp_width_ratio)
+
+        self.img_mod = ModulateWan(self.hidden_size, factor=_MODULATION_FACTOR)
+        self.fused_modulate_img_norm1 = LayerNormScaleShift(
+            self.hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+
+        self.img_attn_qkv = MergedColumnParallelLinear(
+            self.hidden_size,
+            [hidden_size, hidden_size, hidden_size],
+            bias=True,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.img_attn_qkv",
+        )
+        self.img_attn_q_norm = RMSNorm(
+            self.head_dim,
+            eps=1e-6,
+        )
+        self.img_attn_k_norm = RMSNorm(
+            self.head_dim,
+            eps=1e-6,
+        )
+        self.img_attn_proj = RowParallelLinear(
+            self.hidden_size,
+            hidden_size,
+            bias=True,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.img_attn_proj",
+        )
+
+        self.fused_modulate_img_norm2 = LayerNormScaleShift(
+            self.hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+        self.img_mlp = MLP(
+            input_dim=self.hidden_size,
+            mlp_hidden_dim=self.mlp_hidden_dim,
+            act_type=mlp_act_type,
+            quant_config=quant_config,
+            prefix=f"{prefix}.img_mlp",
+        )
+
+        # Text modulation and attention
+        self.txt_mod = ModulateWan(self.hidden_size, factor=_MODULATION_FACTOR)
+        self.fused_modulate_txt_norm1 = LayerNormScaleShift(
+            self.hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+        self.txt_attn_qkv = MergedColumnParallelLinear(
+            self.hidden_size,
+            [self.hidden_size, self.hidden_size, self.hidden_size],
+            bias=True,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.txt_attn_qkv",
+        )
+        self.txt_attn_q_norm = RMSNorm(
+            self.head_dim,
+            eps=1e-6,
+        )
+        self.txt_attn_k_norm = RMSNorm(
+            self.head_dim,
+            eps=1e-6,
+        )
+        self.txt_attn_proj = RowParallelLinear(
+            self.hidden_size,
+            self.hidden_size,
+            bias=True,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.txt_attn_proj",
+        )
+
+        self.fused_modulate_txt_norm2 = LayerNormScaleShift(
+            self.hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+        self.txt_mlp = MLP(
+            input_dim=self.hidden_size,
+            mlp_hidden_dim=self.mlp_hidden_dim,
+            act_type=mlp_act_type,
+            quant_config=quant_config,
+            prefix=f"{prefix}.txt_mlp",
+        )
+        self.attn = USPAttention(
+            num_heads=self.local_heads_num,
+            head_size=self.head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+            softmax_scale=None,
+        )
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        vis_freqs_cis: Optional[torch.Tensor] = None,
+        txt_freqs_cis: Optional[torch.Tensor] = None,
+        num_replicated_suffix: int = 0,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass through multimodal double stream block."""
+        (
+            img_mod1_shift,
+            img_mod1_scale,
+            img_mod1_gate,
+            img_mod2_shift,
+            img_mod2_scale,
+            img_mod2_gate,
+        ) = self.img_mod(vec)
+        (
+            txt_mod1_shift,
+            txt_mod1_scale,
+            txt_mod1_gate,
+            txt_mod2_shift,
+            txt_mod2_scale,
+            txt_mod2_gate,
+        ) = self.txt_mod(vec)
+
+        # Image attention
+        img_modulated = self.fused_modulate_img_norm1(
+            img, shift=img_mod1_shift, scale=img_mod1_scale
+        )
+        img_qkv, _ = self.img_attn_qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(
+            img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.local_heads_num
+        )
+
+        if vis_freqs_cis is None:
+            raise ValueError(
+                "vis_freqs_cis is required for fused QK-Norm + RoPE kernel"
+            )
+        if not (isinstance(vis_freqs_cis, torch.Tensor) and vis_freqs_cis.dim() == 2):
+            raise ValueError("vis_freqs_cis must be a 2D cos_sin_cache tensor")
+        if img_q.dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError(
+                f"Fused QK-Norm + RoPE kernel only supports float16/bfloat16, but got {img_q.dtype}"
+            )
+        img_q = img_q.contiguous()
+        img_k = img_k.contiguous()
+        img_q, img_k = apply_qk_norm_with_optional_rope(
+            q=img_q,
+            k=img_k,
+            q_norm=self.img_attn_q_norm,
+            k_norm=self.img_attn_k_norm,
+            head_dim=img_q.shape[-1],
+            cos_sin_cache=vis_freqs_cis,
+            is_neox=False,
+            allow_inplace=True,
+        )
+        img_q, img_k = img_q.to(img_v), img_k.to(img_v)
+
+        # Text attention
+        txt_modulated = self.fused_modulate_txt_norm1(
+            txt, shift=txt_mod1_shift, scale=txt_mod1_scale
+        )
+        txt_qkv, _ = self.txt_attn_qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(
+            txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.local_heads_num
+        )
+
+        if txt_freqs_cis is not None and not (
+            isinstance(txt_freqs_cis, torch.Tensor) and txt_freqs_cis.dim() == 2
+        ):
+            raise ValueError("txt_freqs_cis must be a 2D cos_sin_cache tensor")
+        txt_q = txt_q.contiguous()
+        txt_k = txt_k.contiguous()
+        txt_q, txt_k = apply_qk_norm_with_optional_rope(
+            q=txt_q,
+            k=txt_k,
+            q_norm=self.txt_attn_q_norm,
+            k_norm=self.txt_attn_k_norm,
+            head_dim=txt_q.shape[-1],
+            cos_sin_cache=txt_freqs_cis,
+            is_neox=False,
+            allow_inplace=True,
+        )
+        txt_q, txt_k = txt_q.to(txt_v), txt_k.to(txt_v)
+
+        # Attention
+        joint_query = torch.cat([img_q, txt_q], dim=1)
+        joint_key = torch.cat([img_k, txt_k], dim=1)
+        joint_value = torch.cat([img_v, txt_v], dim=1)
+        attn = self.attn(
+            joint_query,
+            joint_key,
+            joint_value,
+            num_replicated_suffix=num_replicated_suffix,
+        )
+        attn = attn.flatten(2, 3)
+        img_attn, txt_attn = (
+            attn[:, : img.shape[1]],
+            attn[:, img.shape[1] :],
+        )
+
+        img = fused_add_gate(img, self.img_attn_proj(img_attn)[0], img_mod1_gate)
+        img = fused_add_gate(
+            img,
+            self.img_mlp(
+                self.fused_modulate_img_norm2(
+                    img, shift=img_mod2_shift, scale=img_mod2_scale
+                )
+            ),
+            img_mod2_gate,
+        )
+
+        # Text blocks
+        txt = fused_add_gate(txt, self.txt_attn_proj(txt_attn)[0], txt_mod1_gate)
+        txt = fused_add_gate(
+            txt,
+            self.txt_mlp(
+                self.fused_modulate_txt_norm2(
+                    txt, shift=txt_mod2_shift, scale=txt_mod2_scale
+                )
+            ),
+            txt_mod2_gate,
+        )
+
+        return img, txt
+
+
+class JoyTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """
+    JoyImage Transformer 3D Model for image generation.
+
+    """
+
+    _supports_gradient_checkpointing = True
+    _fsdp_shard_conditions = JoyImageDiTConfig()._fsdp_shard_conditions
+    _compile_conditions = JoyImageDiTConfig()._compile_conditions
+    _supported_attention_backends = JoyImageDiTConfig()._supported_attention_backends
+    param_names_mapping = JoyImageDiTConfig().param_names_mapping
+    reverse_param_names_mapping = JoyImageDiTConfig().reverse_param_names_mapping
+    lora_param_names_mapping = JoyImageDiTConfig().lora_param_names_mapping
+
+    def __init__(
+        self,
+        config: JoyImageDiTConfig,
+        hf_config: dict[str, Any],
+        quant_config: Optional[QuantizationConfig] = None,
+    ) -> None:
+        super().__init__(
+            config=config,
+            hf_config=hf_config,
+        )
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels or config.in_channels
+        self.patch_size = config.patch_size
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.rope_dim_list = config.rope_dim_list
+        self.mm_double_blocks_depth = config.mm_double_blocks_depth
+        self.rope_theta = config.rope_theta
+        self.quant_config = quant_config
+        self.num_channels_latents = self.out_channels
+
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"Hidden size {self.hidden_size} must be divisible by num_attention_heads {self.num_attention_heads}"
+            )
+
+        # Image projection (patch embedding)
+        self.img_in = nn.Conv3d(
+            self.in_channels,
+            self.hidden_size,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+
+        # Condition embedding
+        self.condition_embedder = WanTimeTextImageEmbedding(
+            dim=self.hidden_size,
+            time_freq_dim=config.freq_dim,
+            text_embed_dim=config.text_states_dim,
+        )
+
+        # Double blocks (DiT layers)
+        self.double_blocks = nn.ModuleList(
+            [
+                MMDoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_attention_heads,
+                    mlp_width_ratio=config.mlp_width_ratio,
+                    supported_attention_backends=self._supported_attention_backends,
+                    quant_config=quant_config,
+                    prefix=f"{config.prefix}.double_blocks.{i}",
+                )
+                for i in range(self.mm_double_blocks_depth)
+            ]
+        )
+        # Layerwise offload expects ModuleList names here.
+        self.layer_names = ["double_blocks"]
+
+        # Output norm & projection
+        self.norm_out = nn.LayerNorm(
+            self.hidden_size, elementwise_affine=False, eps=1e-6
+        )
+        self.proj_out = ReplicatedLinear(
+            self.hidden_size,
+            self.out_channels * math.prod(self.patch_size),
+            quant_config=quant_config,
+            prefix="proj_out",
+        )
+        self.__post_init__()
+
+        self.sp_size = get_sp_world_size()
+        self.rotary_emb = NDRotaryEmbedding(
+            rope_dim_list=config.rope_dim_list,
+            rope_theta=config.rope_theta,
+            dtype=torch.float32,
+        )
+
+    @lru_cache(maxsize=1)
+    def _compute_rope_for_local_shard(
+        self,
+        local_len: int,
+        rank: int,
+        vae_image_sizes: tuple[tuple[int, int, int], ...],
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        token_start = rank * local_len
+        token_indices = torch.arange(
+            token_start,
+            token_start + local_len,
+            device=device,
+            dtype=torch.long,
+        )
+        positions = torch.zeros(local_len, 3, device=device, dtype=torch.long)
+
+        cumsum = 0
+        current_t_offset = 0
+        for t, h, w in vae_image_sizes:
+            item_size = t * h * w
+            mask = (token_indices >= cumsum) & (token_indices < cumsum + item_size)
+            if mask.any():
+                local_idx = token_indices[mask] - cumsum
+                frame_stride = h * w
+                positions[mask, 0] = local_idx // frame_stride + current_t_offset
+                positions[mask, 1] = (local_idx % frame_stride) // w
+                positions[mask, 2] = local_idx % w
+            cumsum += item_size
+            current_t_offset += t
+
+        return self.rotary_emb.forward_uncached(positions)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_mask: torch.Tensor | list[torch.Tensor] | None = None,
+        vis_freqs_cis: torch.Tensor | None = None,
+        txt_freqs_cis: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Forward pass through JoyImage Transformer."""
+        forward_batch = get_forward_context().forward_batch
+        sequence_shard_enabled = (
+            forward_batch is not None
+            and getattr(forward_batch, "enable_sequence_shard", False)
+            and self.sp_size > 1
+        )
+
+        batch_size = hidden_states.shape[0]
+
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        if isinstance(encoder_hidden_states_mask, list):
+            encoder_hidden_states_mask = encoder_hidden_states_mask[0]
+
+        cond_batch = int(encoder_hidden_states.shape[0])
+        if cond_batch != int(batch_size):
+            if cond_batch <= 0 or int(batch_size) % cond_batch != 0:
+                raise ValueError(
+                    "JoyImage conditioning batch mismatch: "
+                    f"hidden_states batch={batch_size}, "
+                    f"encoder_hidden_states batch={cond_batch}."
+                )
+            repeat_factor = int(batch_size) // cond_batch
+            encoder_hidden_states = encoder_hidden_states.repeat_interleave(
+                repeat_factor, dim=0
+            )
+            if encoder_hidden_states_mask is not None:
+                encoder_hidden_states_mask = (
+                    encoder_hidden_states_mask.repeat_interleave(repeat_factor, dim=0)
+                )
+
+        # Prepare img
+        x = rearrange(hidden_states, "b n c p1 p2 p3 -> (b n) c p1 p2 p3")
+        x = self.img_in(x)
+        img = rearrange(x, "(b n) d 1 1 1 -> b n d", b=batch_size)
+
+        seq_len_orig = img.shape[1]
+        seq_shard_pad = 0
+        if sequence_shard_enabled:
+            if seq_len_orig % self.sp_size != 0:
+                seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size)
+                pad = torch.zeros(
+                    (batch_size, seq_shard_pad, img.shape[2]),
+                    dtype=img.dtype,
+                    device=img.device,
+                )
+                img = torch.cat([img, pad], dim=1)
+            sp_rank = get_sp_group().rank_in_group
+            local_seq_len = img.shape[1] // self.sp_size
+            img = img.view(batch_size, self.sp_size, local_seq_len, img.shape[2])[
+                :, sp_rank, :, :
+            ].contiguous()
+
+        # Compute rope in model for all SP modes
+        if forward_batch is not None and forward_batch.vae_image_sizes is not None:
+            vae_image_sizes = tuple(tuple(s) for s in forward_batch.vae_image_sizes)
+            local_len = img.shape[1]
+            rank = get_sp_group().rank_in_group if self.sp_size > 1 else 0
+            freqs_cos, freqs_sin = self._compute_rope_for_local_shard(
+                local_len,
+                rank,
+                vae_image_sizes,
+                img.device,
+            )
+            vis_freqs_cis = torch.cat(
+                [
+                    freqs_cos.to(dtype=torch.float32).contiguous(),
+                    freqs_sin.to(dtype=torch.float32).contiguous(),
+                ],
+                dim=-1,
+            )
+
+        _, vec, txt, _ = self.condition_embedder(timestep, encoder_hidden_states)
+        if vec.shape[-1] > self.hidden_size:
+            vec = vec.unflatten(1, (_MODULATION_FACTOR, -1))
+
+        txt_suffix_len = txt.shape[1] if sequence_shard_enabled else 0
+
+        # Pass through DiT blocks
+        for block in self.double_blocks:
+            img, txt = block(
+                img,
+                txt,
+                vec,
+                vis_freqs_cis,
+                txt_freqs_cis,
+                num_replicated_suffix=txt_suffix_len,
+            )
+
+        if sequence_shard_enabled:
+            img = img.contiguous()
+            img = sequence_model_parallel_all_gather(img, dim=1)
+            if seq_shard_pad > 0:
+                img = img[:, :seq_len_orig, :]
+
+        img, _ = self.proj_out(self.norm_out(img))
+
+        # Restore patch layout expected by downstream latent decoding.
+        img = rearrange(
+            img,
+            "b n (pt ph pw c) -> b n c pt ph pw",
+            pt=self.patch_size[0],
+            ph=self.patch_size[1],
+            pw=self.patch_size[2],
+            c=self.out_channels,
+        )
+
+        return img
+
+
+class JoyImageEditTransformer3DModel(JoyTransformer3DModel):
+    """Backward-compatible alias for JoyImageEdit model configs."""
+
+    pass
+
+
+EntryClass = [JoyTransformer3DModel, JoyImageEditTransformer3DModel]

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

@@ -0,0 +1,634 @@
+"""Krea-2 (K2) single-stream MMDiT.
+
+Text and image tokens are concatenated into a single joint-attention stream. The
+model uses GQA attention with a sigmoid output gate, 6-way shared adaLN
+modulation, a text-fusion transformer that fuses the selected text-encoder
+hidden-state layers into one, and interleaved 3-axis RoPE. Module and parameter
+names follow the released K2 checkpoint, so weights load without remapping.
+"""
+
+import math
+import os
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+
+from sglang.multimodal_gen.configs.models.dits.krea2 import Krea2DitConfig
+from sglang.multimodal_gen.runtime.distributed import (
+    get_sp_world_size,
+    get_tp_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.attention.layer import build_varlen_mask_meta
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# --------------------------------------------------------------------------- #
+# Functional helpers
+# --------------------------------------------------------------------------- #
+def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> Tensor:
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / ((theta * ntk) ** scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
+    )
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def ropeapply(xq: Tensor, xk: Tensor, freqs: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    freqs = freqs[:, None, :, :, :]
+    xq_ = freqs[..., 0] * xq_[..., 0] + freqs[..., 1] * xq_[..., 1]
+    xk_ = freqs[..., 0] * xk_[..., 0] + freqs[..., 1] * xk_[..., 1]
+    return xq_.reshape(*xq.shape).to(xq.dtype), xk_.reshape(*xk.shape).to(xk.dtype)
+
+
+def _fused_qknorm_rope_enabled() -> bool:
+    return os.getenv("SGLANG_ENABLE_FUSED_QKNORM_ROPE", "1").lower() not in (
+        "0",
+        "false",
+        "off",
+        "no",
+    )
+
+
+def _can_use_fused_qknorm_rope(head_dim: int, dtype: torch.dtype) -> bool:
+    from sglang.jit_kernel.diffusion.qknorm_rope import (
+        can_use_fused_inplace_qknorm_rope,
+    )
+
+    return can_use_fused_inplace_qknorm_rope(head_dim, head_dim, False, dtype)
+
+
+def _qknorm_rope_cos_sin_cache(freqs: Tensor) -> Tensor:
+    """``[num_tokens, head_dim]`` cos|sin cache for the fused QKNorm+RoPE kernel.
+
+    K2's ``rope`` packs each token's rotation as ``[[cos, -sin], [sin, cos]]`` in a
+    ``[B, N, head_dim//2, 2, 2]`` tensor; the kernel wants the per-token cosines then
+    sines concatenated. Positions come from the image grid (batch-invariant), so the
+    first batch row is representative.
+    """
+    return torch.cat([freqs[0, :, :, 0, 0], freqs[0, :, :, 1, 0]], dim=-1).float()
+
+
+def temb(
+    t: Tensor,
+    dim: int,
+    period: float = 1e4,
+    tfactor: float = 1e3,
+    device: torch.device = None,
+    dtype: torch.dtype = None,
+) -> Tensor:
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(period)
+        * torch.arange(half, dtype=torch.float32, device=device)
+        / half
+    )
+    args = (t.float() * tfactor)[:, None, None] * freqs
+    sin, cos = torch.sin(args), torch.cos(args)
+    return torch.cat((cos, sin), dim=-1).to(dtype=dtype)
+
+
+def norm_scale_shift(
+    x: Tensor, weight: Tensor, scale: Tensor, shift: Tensor, eps: float
+) -> Tensor:
+    """Fused RMSNorm + modulation: ``rms_norm(x) * weight * (1 + scale) + shift``.
+
+    ``weight`` is the effective RMSNorm weight (K2 stores ``scale``, so callers
+    pass ``scale + 1``), kept off the checkpoint so the identity load is unaffected.
+    """
+    if x.is_cuda and x.shape[-1] % 256 == 0:
+        from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import (
+            fused_norm_scale_shift,
+        )
+
+        return fused_norm_scale_shift(
+            x.contiguous(),
+            weight.contiguous(),
+            None,
+            scale.contiguous(),
+            shift.contiguous(),
+            "rms",
+            eps,
+        )
+    normed = F.rms_norm(x.float(), (x.shape[-1],), weight=weight.float(), eps=eps)
+    return (normed.to(x.dtype) * (1 + scale) + shift).to(x.dtype)
+
+
+# --------------------------------------------------------------------------- #
+# Submodules
+# --------------------------------------------------------------------------- #
+class TimeEmbed(nn.Module):
+    """Timestep embedding MLP: linear_1 -> gelu(tanh) -> linear_2."""
+
+    def __init__(self, in_dim: int, dim: int):
+        super().__init__()
+        self.linear_1 = nn.Linear(in_dim, dim)
+        self.linear_2 = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.linear_2(F.gelu(self.linear_1(x), approximate="tanh"))
+
+
+class TxtIn(nn.Module):
+    """Text-context projection: rms-norm -> linear_1 -> gelu(tanh) -> linear_2."""
+
+    def __init__(self, txt_dim: int, dim: int):
+        super().__init__()
+        self.norm = RMSNorm(txt_dim)
+        self.linear_1 = nn.Linear(txt_dim, dim)
+        self.linear_2 = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.linear_2(F.gelu(self.linear_1(self.norm(x)), approximate="tanh"))
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, dim, axdims: list[int], theta: float = 1e2, ntk: float = 1.0):
+        super().__init__()
+        self.axdims = axdims
+        self.theta = theta
+        self.ntk = ntk
+
+    def forward(self, pos: Tensor) -> Tensor:
+        return torch.cat(
+            [
+                rope(pos[..., i], d, self.theta, self.ntk)
+                for i, d in enumerate(self.axdims)
+            ],
+            dim=-3,
+        )
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm with effective scale ``weight + 1`` (``weight`` initialized to 0),
+    computed in fp32. The parameter is named ``weight`` to match the released
+    checkpoint; the ``+ 1`` is applied in the forward."""
+
+    def __init__(self, features: int, eps: float = 1e-05, device: torch.device = None):
+        super().__init__()
+        self.features = features
+        self.eps = eps
+        self.weight = nn.Parameter(
+            torch.zeros(features, device=device, dtype=torch.float32)
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        t, dtype = x.float(), x.dtype
+        t = F.rms_norm(
+            t, (self.features,), eps=self.eps, weight=(self.weight.float() + 1.0)
+        )
+        return t.to(dtype)
+
+
+class SwiGLU(nn.Module):
+    def __init__(
+        self, features: int, multiplier: int, bias: bool = False, multiple: int = 128
+    ):
+        super().__init__()
+        mlpdim = int(2 * features / 3) * multiplier
+        mlpdim = multiple * ((mlpdim + multiple - 1) // multiple)
+        # Tensor-parallel: gate/up shard the hidden dim by column, down all-reduces.
+        self.gate = ColumnParallelLinear(
+            features, mlpdim, bias=bias, gather_output=False
+        )
+        self.up = ColumnParallelLinear(features, mlpdim, bias=bias, gather_output=False)
+        self.down = RowParallelLinear(
+            mlpdim, features, bias=bias, input_is_parallel=True
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        gate, _ = self.gate(x)
+        up, _ = self.up(x)
+        out, _ = self.down(F.silu(gate) * up)
+        return out
+
+
+class Attention(nn.Module):
+    def __init__(self, dim: int, heads: int, kvheads: int = None, bias: bool = False):
+        super().__init__()
+        self.heads = heads
+        self.kvheads = kvheads if kvheads is not None else heads
+        self.headdim = dim // self.heads
+
+        # Tensor-parallel: q/k/v/gate shard heads by column, to_out all-reduces.
+        # Parameter names match the released checkpoint (to_q/to_k/to_v/to_gate,
+        # norm_q/norm_k, to_out.0) so the checkpoint loads with an identity mapping.
+        tp = get_tp_world_size()
+        assert (
+            self.heads % tp == 0 and self.kvheads % tp == 0
+        ), f"heads={self.heads}, kvheads={self.kvheads} must be divisible by tp={tp}"
+        self.local_heads = self.heads // tp
+        self.local_kvheads = self.kvheads // tp
+
+        self.to_q = ColumnParallelLinear(
+            dim, self.headdim * self.heads, bias=bias, gather_output=False
+        )
+        self.to_k = ColumnParallelLinear(
+            dim, self.headdim * self.kvheads, bias=bias, gather_output=False
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, self.headdim * self.kvheads, bias=bias, gather_output=False
+        )
+        self.to_gate = ColumnParallelLinear(dim, dim, bias=bias, gather_output=False)
+        self.norm_q = RMSNorm(self.headdim)
+        self.norm_k = RMSNorm(self.headdim)
+        # to_out is a ModuleList ([linear]) so the param is to_out.0.weight, matching
+        # the diffusers Attention layout in the released checkpoint.
+        self.to_out = nn.ModuleList(
+            [RowParallelLinear(dim, dim, bias=bias, input_is_parallel=True)]
+        )
+        # Native GQA flash via the platform backend; parameterless.
+        self.attn = USPAttention(
+            num_heads=self.local_heads,
+            head_size=self.headdim,
+            num_kv_heads=self.local_kvheads,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+        )
+
+    def forward(
+        self,
+        qkv: Tensor,
+        freqs: Tensor | None = None,
+        key_mask: Tensor | None = None,
+        mask_meta: dict | None = None,
+        num_replicated_prefix: int = 0,
+        skip_sequence_parallel: bool = False,
+    ) -> Tensor:
+        q, _ = self.to_q(qkv)
+        k, _ = self.to_k(qkv)
+        v, _ = self.to_v(qkv)
+        gate, _ = self.to_gate(qkv)
+
+        hd = self.headdim
+        # Fast path: fuse RMSNorm(q), RMSNorm(k) and RoPE into one in-place kernel on
+        # the [B, S, H, D] layout USPAttention consumes (also skips the [B, H, L, D]
+        # transpose round-trip the eager path needs). Eager fallback below preserves
+        # parity off CUDA / for unsupported dtypes.
+        if (
+            freqs is not None
+            and q.is_cuda
+            and q.dtype in (torch.float16, torch.bfloat16)
+            and _fused_qknorm_rope_enabled()
+            and _can_use_fused_qknorm_rope(hd, q.dtype)
+        ):
+            from sglang.jit_kernel.diffusion.qknorm_rope import (
+                fused_inplace_qknorm_rope,
+            )
+
+            b, s = qkv.shape[0], qkv.shape[1]
+            q = q.view(b, s, self.local_heads, hd)
+            k = k.view(b, s, self.local_kvheads, hd)
+            v = v.view(b, s, self.local_kvheads, hd)
+            positions = torch.arange(s, device=q.device, dtype=torch.long)
+            if b > 1:
+                positions = positions.repeat(b)
+            fused_inplace_qknorm_rope(
+                q.reshape(-1, self.local_heads, hd),
+                k.reshape(-1, self.local_kvheads, hd),
+                (self.norm_q.weight.float() + 1.0).to(q.dtype),
+                (self.norm_k.weight.float() + 1.0).to(k.dtype),
+                _qknorm_rope_cos_sin_cache(freqs),
+                positions,
+                is_neox=False,
+                eps=self.norm_q.eps,
+                head_dim=hd,
+                rope_dim=hd,
+            )
+            out = self.attn(
+                q,
+                k,
+                v,
+                attn_mask=key_mask,
+                attn_mask_meta=mask_meta,
+                num_replicated_prefix=num_replicated_prefix,
+                skip_sequence_parallel_override=skip_sequence_parallel,
+            ).flatten(2)
+        else:
+            q, k, v = (
+                rearrange(q, "B L (H D) -> B H L D", H=self.local_heads),
+                rearrange(k, "B L (H D) -> B H L D", H=self.local_kvheads),
+                rearrange(v, "B L (H D) -> B H L D", H=self.local_kvheads),
+            )
+            q, k = self.norm_q(q), self.norm_k(k)
+            if freqs is not None:
+                q, k = ropeapply(q, k, freqs)
+            # USPAttention expects [B, S, H, D]; a [B, S] key mask + varlen metadata
+            # routes a ragged batch through the FA varlen fast path, else maskless.
+            out = self.attn(
+                q.transpose(1, 2).contiguous(),
+                k.transpose(1, 2).contiguous(),
+                v.transpose(1, 2).contiguous(),
+                attn_mask=key_mask,
+                attn_mask_meta=mask_meta,
+                num_replicated_prefix=num_replicated_prefix,
+                skip_sequence_parallel_override=skip_sequence_parallel,
+            ).flatten(2)
+        out, _ = self.to_out[0](out * F.sigmoid(gate))
+        return out
+
+
+class LastLayer(nn.Module):
+    def __init__(self, features: int, patch: int, channels: int):
+        super().__init__()
+        self.norm = RMSNorm(features)
+        self.linear = nn.Linear(features, patch * patch * channels, bias=True)
+        self.scale_shift_table = nn.Parameter(torch.zeros(2, features))
+
+    def forward(self, x: Tensor, tvec: Tensor) -> Tensor:
+        mod = tvec + rearrange(self.scale_shift_table, "two d -> 1 two d")
+        scale, shift = mod.chunk(2, dim=1)
+        x = norm_scale_shift(x, self.norm.weight + 1, scale, shift, self.norm.eps)
+        x = self.linear(x)
+        return x
+
+
+class TextFusionBlock(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = None,
+    ):
+        super().__init__()
+        self.norm1 = RMSNorm(features)
+        self.norm2 = RMSNorm(features)
+        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.ff = SwiGLU(features, multiplier, bias)
+
+    def forward(
+        self,
+        x: Tensor,
+        key_mask: Tensor | None = None,
+        mask_meta: dict | None = None,
+    ) -> Tensor:
+        # Text-fusion runs on the full replicated text, so skip the SP all-to-all.
+        x = x + self.attn(
+            self.norm1(x),
+            key_mask=key_mask,
+            mask_meta=mask_meta,
+            skip_sequence_parallel=True,
+        )
+        x = x + self.ff(self.norm2(x))
+        return x
+
+
+class TextFusionTransformer(nn.Module):
+    """Fuses `num_txt_layers` selected encoder hidden-state layers into one.
+
+    Depth is fixed at 2 layerwise + 2 refiner blocks; `num_txt_layers` is the
+    projector input width (the layer axis), NOT the transformer depth.
+    """
+
+    def __init__(
+        self,
+        num_txt_layers: int,
+        txt_dim: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = None,
+    ):
+        super().__init__()
+        self.layerwise_blocks = nn.ModuleList(
+            [
+                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                for _ in range(2)
+            ]
+        )
+        self.projector = nn.Linear(num_txt_layers, 1, bias=False)
+        self.refiner_blocks = nn.ModuleList(
+            [
+                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                for _ in range(2)
+            ]
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        key_mask: Tensor | None = None,
+        mask_meta: dict | None = None,
+    ) -> Tensor:
+        b, l, n, d = x.shape
+        x = x.reshape(b * l, n, d)
+        for block in self.layerwise_blocks:
+            x = block(x.contiguous())
+        x = rearrange(x, "(b l) n d -> b l d n", b=b, l=l)
+        x = self.projector(x)
+        x = x.squeeze(-1)
+        for block in self.refiner_blocks:
+            x = block(x, key_mask=key_mask, mask_meta=mask_meta)
+        return x
+
+
+class SingleStreamBlock(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = None,
+    ):
+        super().__init__()
+        # (6, features) modulation table added to the timestep projection (AdaLN-single),
+        # stored directly on the block to match the released checkpoint.
+        self.scale_shift_table = nn.Parameter(torch.zeros(6, features))
+        self.norm1 = RMSNorm(features)
+        self.norm2 = RMSNorm(features)
+        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.ff = SwiGLU(features, multiplier, bias)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        vec: Tensor,
+        freqs: Tensor,
+        key_mask: Tensor | None = None,
+        mask_meta: dict | None = None,
+        num_replicated_prefix: int = 0,
+    ) -> Tensor:
+        mod = vec + self.scale_shift_table.reshape(-1)
+        prescale, preshift, pregate, postscale, postshift, postgate = mod.chunk(
+            6, dim=-1
+        )
+        hidden_states = hidden_states + pregate * self.attn(
+            norm_scale_shift(
+                hidden_states,
+                self.norm1.weight + 1,
+                prescale,
+                preshift,
+                self.norm1.eps,
+            ),
+            freqs,
+            key_mask,
+            mask_meta,
+            num_replicated_prefix=num_replicated_prefix,
+        )
+        hidden_states = hidden_states + postgate * self.ff(
+            norm_scale_shift(
+                hidden_states,
+                self.norm2.weight + 1,
+                postscale,
+                postshift,
+                self.norm2.eps,
+            )
+        )
+        return hidden_states
+
+
+# --------------------------------------------------------------------------- #
+# Top-level model
+# --------------------------------------------------------------------------- #
+class Krea2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """K2 single-stream MMDiT for the SGLang diffusion runtime.
+
+    Attribute names follow the released K2 checkpoint, so weights load with an
+    identity ``param_names_mapping``.
+    """
+
+    _fsdp_shard_conditions = []
+    _compile_conditions = []
+    param_names_mapping = Krea2DitConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = {}
+
+    def __init__(
+        self,
+        config: Krea2DitConfig,
+        hf_config: dict[str, Any],
+        quant_config: Optional[Any] = None,
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        ac = config.arch_config
+        self.arch_config = ac
+
+        self.hidden_size = ac.features
+        self.num_attention_heads = ac.heads
+        self.num_channels_latents = ac.channels
+        self.patch = ac.patch
+        self.channels = ac.channels
+        self.tdim = ac.tdim
+
+        head_dim = ac.features // ac.heads
+        axes = list(ac.axes_dims)
+        assert sum(axes) == head_dim, f"sum(axes)={sum(axes)}, head_dim={head_dim}"
+        assert all(a % 2 == 0 for a in axes), f"axes={axes}"
+
+        self.posemb = PositionalEncoding(ac.features, axes, theta=ac.theta, ntk=1.0)
+        self.img_in = nn.Linear(ac.channels * ac.patch**2, ac.features, bias=True)
+        self.transformer_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(
+                    ac.features, ac.heads, ac.multiplier, ac.bias, ac.kvheads
+                )
+                for _ in range(ac.layers)
+            ]
+        )
+        self.time_embed = TimeEmbed(ac.tdim, ac.features)
+        self.text_fusion = TextFusionTransformer(
+            ac.txtlayers,
+            ac.txtdim,
+            ac.txtheads,
+            ac.multiplier,
+            ac.bias,
+            ac.txtkvheads,
+        )
+        self.txt_in = TxtIn(ac.txtdim, ac.features)
+        self.final_layer = LastLayer(ac.features, ac.patch, ac.channels)
+        # GELU(tanh) is applied in the forward; the linear matches time_mod_proj.weight.
+        self.time_mod_proj = nn.Linear(ac.features, ac.features * 6)
+        self.seq_multiple_of = ac.seq_multiple_of
+        # The 28 single-stream blocks (the ~24GB bulk) are streamed layer-by-layer
+        # under --dit-layerwise-offload, keeping only a small working set resident.
+        self.layer_names = ["transformer_blocks"]
+
+    def _forward_impl(
+        self,
+        img: Tensor,
+        context: Tensor,
+        t: Tensor,
+        pos: Tensor,
+        mask: Tensor | None = None,
+    ) -> Tensor:
+        img = self.img_in(img)
+        t = self.time_embed(temb(t, self.tdim, device=img.device, dtype=img.dtype))
+        tvec = self.time_mod_proj(F.gelu(t, approximate="tanh"))
+
+        # A single or same-prompt batch has no padding, so attention runs maskless
+        # (native-GQA flash). A ragged batch builds varlen metadata from the
+        # key mask and takes the FA varlen path instead.
+        txt_key = txt_meta = joint_key = joint_meta = None
+        if mask is not None and not bool(mask.all()):
+            txt_key = mask[:, : context.shape[1]]
+            txt_meta = build_varlen_mask_meta(txt_key)
+            joint_key = mask
+            joint_meta = build_varlen_mask_meta(mask)
+
+        context = self.text_fusion(context, key_mask=txt_key, mask_meta=txt_meta)
+        context = self.txt_in(context)
+
+        txtlen, imglen = context.shape[1], img.shape[1]
+        combined = torch.cat((context, img), dim=1)
+        freqs = self.posemb(pos)
+
+        # Under SP the image tokens are sharded across ranks while the text prefix
+        # stays replicated; keep the leading txtlen tokens out of the all-to-all.
+        num_replicated_prefix = txtlen if get_sp_world_size() > 1 else 0
+        for block in self.transformer_blocks:
+            combined = block(
+                combined,
+                tvec,
+                freqs,
+                joint_key,
+                joint_meta,
+                num_replicated_prefix=num_replicated_prefix,
+            )
+
+        final = self.final_layer(combined, t)
+        output = final[:, txtlen : txtlen + imglen, :]
+        return output
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        encoder_hidden_states: Tensor,
+        timestep: Tensor,
+        encoder_hidden_states_image=None,
+        guidance=None,
+        pos: Tensor = None,
+        mask: Tensor = None,
+        **kwargs,
+    ) -> Tensor:
+        return self._forward_impl(
+            img=hidden_states,
+            context=encoder_hidden_states,
+            t=timestep,
+            pos=pos,
+            mask=mask,
+        )
+
+
+EntryClass = [Krea2Transformer2DModel]

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

@@ -0,0 +1,269 @@
+# 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.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+
+# 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, LayerwiseOffloadableModuleMixin):
+    _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],
+        quant_config: QuantizationConfig | None = None,
+    ) -> 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
+        separated_timestep = config.separated_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.separated_timestep = separated_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",
+            quant_config=quant_config,
+        )
+        self.time_embedding = MLP(
+            freq_dim, dim, output_dim=dim, act_type="silu", quant_config=quant_config
+        )
+        # Preserve state_dict keys (time_projection.1.weight/bias).
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), ReplicatedLinear(dim, dim * 6, quant_config=quant_config)
+        )
+        self.blocks = nn.ModuleList(
+            [
+                DiTBlock(dim, num_heads, ffn_dim, eps, quant_config=quant_config)
+                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__()
+
+    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)
+        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,
+    ) -> 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)
+        )
+
+        for block in self.blocks:
+            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,595 @@
+# 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 (
+    LayerNormScaleShift,
+    RMSNorm,
+    ScaleResidualLayerNormScaleShift,
+    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.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.platforms import current_platform
+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,
+        quant_config: QuantizationConfig | None = None,
+    ):
+        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, quant_config=quant_config
+        )
+        self.k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.o = RowParallelLinear(
+            dim, dim, bias=True, input_is_parallel=True, quant_config=quant_config
+        )
+        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, attn_mask_meta=None):
+        """
+        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
+            attn_mask_meta: sp_shard tail-pad meta; excludes SP padding from attention
+
+        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; the tail meta keeps
+        # SP padding out of the softmax.
+        out = self.attn(q, k, v, attn_mask_meta=attn_mask_meta)
+        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,
+        quant_config: QuantizationConfig | None = None,
+    ):
+        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, quant_config=quant_config
+        )
+        self.k = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.v = ColumnParallelLinear(
+            dim, dim, bias=True, gather_output=False, quant_config=quant_config
+        )
+        self.o = RowParallelLinear(
+            dim, dim, bias=True, input_is_parallel=True, quant_config=quant_config
+        )
+        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 MulAdd(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, gate, residual):
+        return residual + gate * x
+
+
+class DiTBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        ffn_dim: int,
+        eps: float = 1e-6,
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+
+        self.self_attn = SelfAttention(dim, num_heads, eps, quant_config=quant_config)
+        self.cross_attn = CrossAttention(dim, num_heads, eps, quant_config=quant_config)
+        self.norm1 = LayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
+        )
+        self.self_attn_norm = nn.LayerNorm(dim, eps=eps)
+        # Fused: residual + 1 * cross_attn_out → layernorm + scale/shift
+        # Replaces the old norm2 (LayerNormScaleShift) + residual add for cross-attention
+        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
+        )
+        self.ffn = MLP(
+            dim,
+            ffn_dim,
+            output_dim=dim,
+            act_type="gelu_pytorch_tanh",
+            quant_config=quant_config,
+        )
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.mlp_residual = MulAdd()
+
+    def forward(self, x, context, t_mod, freqs, attn_mask_meta=None):
+        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),
+            )
+        orig_dtype = x.dtype
+        # 1. Self-attention, fuse:
+        # - layernorm(x) * (1 + scale_msa) + shift_msa
+        input_x = self.norm1(x, shift_msa, scale_msa)
+        # 2. torch.compile may fuse mlp_residual and self_attn_norm
+        x = self.mlp_residual(
+            self.self_attn(input_x, freqs, attn_mask_meta=attn_mask_meta), gate_msa, x
+        )
+        norm_x = self.self_attn_norm(x)
+        # 3. Cross-attention, fuse:
+        # - x = x + 1 * cross_output
+        # - input_x = layernorm(x) * (1 + scale_mlp) + shift_mlp
+        cross_output = self.cross_attn(norm_x, context)
+        input_x, x = self.cross_attn_residual_norm(
+            x, cross_output, 1, shift_mlp, scale_mlp
+        )
+        # 4. Feed-forward
+        x = self.mlp_residual(self.ffn(input_x), gate_mlp, x)
+        x = x.to(orig_dtype)
+        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 = LayerNormScaleShift(
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
+        )
+        # 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, shift.squeeze(2), scale.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, shift, scale))
+        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, LayerwiseOffloadableModuleMixin):
+    _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],
+        quant_config: QuantizationConfig | None = None,
+    ) -> 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
+        separated_timestep = config.separated_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.separated_timestep = separated_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",
+            quant_config=quant_config,
+        )
+        self.time_embedding = MLP(
+            freq_dim, dim, output_dim=dim, act_type="silu", quant_config=quant_config
+        )
+        # Preserve state_dict keys (time_projection.1.weight/bias).
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), ReplicatedLinear(dim, dim * 6, quant_config=quant_config)
+        )
+        self.blocks = nn.ModuleList(
+            [
+                DiTBlock(dim, num_heads, ffn_dim, eps, quant_config=quant_config)
+                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__()
+
+    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
+    ):
+        if current_platform.is_npu:
+            # torch.channels_last_3d is not supported on NPU
+            x = x.contiguous()
+        else:
+            # NOTE(dhyu): avoid slow_conv
+            x = x.contiguous(memory_format=torch.channels_last_3d)
+        x = self.patch_embedding(x)
+        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,
+    ) -> 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)
+        )
+
+        for block in self.blocks:
+            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/dits/sana.py

@@ -0,0 +1,394 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding
+
+from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.linear import MergedColumnParallelLinear
+from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SanaCombinedTimestepSizeEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.time_proj = Timesteps(
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim
+        )
+
+    def forward(self, timestep, hidden_dtype=None):
+        timesteps_proj = self.time_proj(timestep)
+        if hidden_dtype is not None:
+            timesteps_proj = timesteps_proj.to(dtype=hidden_dtype)
+        timesteps_emb = self.timestep_embedder(timesteps_proj)
+        return timesteps_emb
+
+
+class SanaAdaLayerNormSingle(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.emb = SanaCombinedTimestepSizeEmbeddings(embedding_dim)
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)
+
+    def forward(self, timestep, hidden_dtype=None):
+        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
+        out = self.linear(self.silu(embedded_timestep))
+        return out, embedded_timestep
+
+
+class SanaModulatedNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+    def forward(self, x, temb, scale_shift_table):
+        x = self.norm(x)
+        shift, scale = (scale_shift_table[None] + temb[:, None]).chunk(2, dim=1)
+        x = x * (1 + scale) + shift
+        return x
+
+
+class GLUMBConv(nn.Module):
+    """Gated Linear Unit with Multi-Branch Convolution."""
+
+    def __init__(self, in_channels, out_channels, expand_ratio=2.5):
+        super().__init__()
+        hidden_channels = int(expand_ratio * in_channels)
+        self.nonlinearity = nn.SiLU()
+        self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0)
+        self.conv_depth = nn.Conv2d(
+            hidden_channels * 2,
+            hidden_channels * 2,
+            3,
+            1,
+            1,
+            groups=hidden_channels * 2,
+        )
+        self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False)
+
+    def forward(self, hidden_states):
+        hidden_states = self.conv_inverted(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv_depth(hidden_states)
+        hidden_states, gate = torch.chunk(hidden_states, 2, dim=1)
+        hidden_states = hidden_states * self.nonlinearity(gate)
+        hidden_states = self.conv_point(hidden_states)
+        return hidden_states
+
+
+class SanaLinearAttention(nn.Module):
+    """Linear attention with O(N*D^2) complexity instead of O(N^2*D)."""
+
+    def __init__(self, query_dim, num_heads, head_dim, bias=False):
+        super().__init__()
+        inner_dim = num_heads * head_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        self.inner_dim = inner_dim
+        # Self-attention q/k/v share the same input -> one packed GEMM.
+        self.to_qkv = MergedColumnParallelLinear(
+            query_dim, [inner_dim, inner_dim, inner_dim], bias=bias, gather_output=True
+        )
+        self.to_out = nn.ModuleList(
+            [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()]
+        )
+
+    def forward(self, hidden_states):
+        B, S, _ = hidden_states.shape
+
+        qkv, _ = self.to_qkv(hidden_states)
+        query, key, value = qkv.split(
+            [self.inner_dim, self.inner_dim, self.inner_dim], dim=-1
+        )
+
+        query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        key = key.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        value = value.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        query = F.relu(query)
+        key = F.relu(key)
+
+        kv = torch.matmul(key.transpose(-2, -1), value)  # (B, H, D, D)
+        qkv = torch.matmul(query, kv)  # (B, H, S, D)
+
+        key_sum = key.sum(dim=-2, keepdim=True)  # (B, H, 1, D)
+        normalizer = torch.matmul(query, key_sum.transpose(-2, -1)).clamp(min=1e-6)
+        hidden_states = qkv / normalizer
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1)
+        hidden_states = self.to_out[0](hidden_states)
+        return hidden_states
+
+
+class SanaCrossAttention(nn.Module):
+    def __init__(self, query_dim, cross_attention_dim, num_heads, head_dim, bias=False):
+        super().__init__()
+        inner_dim = num_heads * head_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        self.inner_dim = inner_dim
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)
+        # k/v share the (step-invariant) encoder input -> one packed GEMM.
+        self.to_kv = MergedColumnParallelLinear(
+            cross_attention_dim, [inner_dim, inner_dim], bias=bias, gather_output=True
+        )
+        self.to_out = nn.ModuleList(
+            [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()]
+        )
+
+    def forward(
+        self, hidden_states, encoder_hidden_states, encoder_attention_mask=None
+    ):
+        B, S, _ = hidden_states.shape
+        T = encoder_hidden_states.shape[1]
+
+        query = self.to_q(hidden_states)
+        kv, _ = self.to_kv(encoder_hidden_states)
+        key, value = kv.split([self.inner_dim, self.inner_dim], dim=-1)
+
+        query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        key = key.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        value = value.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attn_mask = None
+        if encoder_attention_mask is not None:
+            attn_mask = encoder_attention_mask.bool()
+            attn_mask = attn_mask[:, None, None, :].expand(B, self.num_heads, S, T)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attn_mask
+        )
+        hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1)
+        hidden_states = self.to_out[0](hidden_states)
+        return hidden_states
+
+
+class SanaTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_attention_heads,
+        attention_head_dim,
+        num_cross_attention_heads,
+        cross_attention_head_dim,
+        cross_attention_dim,
+        mlp_ratio,
+        norm_eps,
+        attention_bias=False,
+    ):
+        super().__init__()
+
+        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps)
+        self.attn1 = SanaLinearAttention(
+            query_dim=dim,
+            num_heads=num_attention_heads,
+            head_dim=attention_head_dim,
+            bias=attention_bias,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps)
+        self.attn2 = SanaCrossAttention(
+            query_dim=dim,
+            cross_attention_dim=cross_attention_dim,
+            num_heads=num_cross_attention_heads,
+            head_dim=cross_attention_head_dim,
+            bias=True,
+        )
+
+        self.ff = GLUMBConv(in_channels=dim, out_channels=dim, expand_ratio=mlp_ratio)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states,
+        timestep,
+        height,
+        width,
+        encoder_attention_mask=None,
+    ):
+        batch_size = hidden_states.shape[0]
+
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+        ).chunk(6, dim=1)
+
+        norm_hidden = self.norm1(hidden_states)
+        norm_hidden = norm_hidden * (1 + scale_msa) + shift_msa
+        attn_output = self.attn1(norm_hidden)
+        hidden_states = hidden_states + gate_msa * attn_output
+
+        attn_output = self.attn2(
+            hidden_states, encoder_hidden_states, encoder_attention_mask
+        )
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden = self.norm2(hidden_states)
+        norm_hidden = norm_hidden * (1 + scale_mlp) + shift_mlp
+        norm_hidden = norm_hidden.unflatten(1, (height, width)).permute(0, 3, 1, 2)
+        ff_output = self.ff(norm_hidden)
+        ff_output = ff_output.flatten(2, 3).permute(0, 2, 1)
+        hidden_states = hidden_states + gate_mlp * ff_output
+
+        return hidden_states
+
+
+class SanaTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+
+    _fsdp_shard_conditions = [
+        lambda n, m: isinstance(m, SanaTransformerBlock),
+    ]
+    _compile_conditions = [
+        lambda n, m: isinstance(m, SanaTransformerBlock),
+    ]
+    param_names_mapping = SanaConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = {}
+
+    def __init__(self, config: SanaConfig, hf_config=None, **kwargs):
+        super().__init__(config, hf_config=hf_config or {}, **kwargs)
+
+        arch = config.arch_config
+        self.out_channels = arch.out_channels
+        self.patch_size = arch.patch_size
+        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
+
+        self.hidden_size = self.inner_dim
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.num_channels_latents
+
+        self.patch_embed = nn.ModuleDict(
+            {
+                "proj": nn.Conv2d(
+                    arch.in_channels,
+                    self.inner_dim,
+                    kernel_size=arch.patch_size,
+                    stride=arch.patch_size,
+                    bias=True,
+                ),
+            }
+        )
+        self.time_embed = SanaAdaLayerNormSingle(self.inner_dim)
+        self.caption_projection = PixArtAlphaTextProjection(
+            in_features=arch.caption_channels,
+            hidden_size=self.inner_dim,
+        )
+
+        self.caption_norm = RMSNorm(self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                SanaTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    num_cross_attention_heads=arch.num_cross_attention_heads,
+                    cross_attention_head_dim=arch.cross_attention_head_dim,
+                    cross_attention_dim=arch.cross_attention_dim,
+                    mlp_ratio=arch.mlp_ratio,
+                    norm_eps=arch.norm_eps,
+                    attention_bias=False,
+                )
+                for _ in range(arch.num_layers)
+            ]
+        )
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, self.inner_dim) / self.inner_dim**0.5
+        )
+
+        self.norm_out = SanaModulatedNorm(self.inner_dim, eps=arch.norm_eps)
+
+        self.proj_out = nn.Linear(
+            self.inner_dim,
+            arch.patch_size * arch.patch_size * self.out_channels,
+            bias=True,
+        )
+
+        self.layer_names = ["transformer_blocks"]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        guidance: torch.Tensor = None,
+        encoder_attention_mask: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+
+        # Input validation - fail fast
+        if encoder_hidden_states is None:
+            raise ValueError("SANA forward pass requires encoder_hidden_states")
+
+        batch_size, channels, height, width = hidden_states.shape
+        p = self.patch_size
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        hidden_states = self.patch_embed["proj"](hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        timestep_emb, embedded_timestep = self.time_embed(
+            timestep, hidden_dtype=hidden_states.dtype
+        )
+
+        if isinstance(encoder_attention_mask, (list, tuple)):
+            encoder_attention_mask = encoder_attention_mask[0]
+
+        encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+        if encoder_hidden_states.shape[0] != batch_size:
+            encoder_hidden_states = encoder_hidden_states.expand(
+                batch_size, -1, -1
+            ).contiguous()
+        encoder_hidden_states = encoder_hidden_states.view(
+            batch_size, -1, hidden_states.shape[-1]
+        )
+        encoder_hidden_states = self.caption_norm(encoder_hidden_states)
+
+        if (
+            encoder_attention_mask is not None
+            and encoder_attention_mask.shape[0] != batch_size
+        ):
+            encoder_attention_mask = encoder_attention_mask.expand(
+                batch_size, -1
+            ).contiguous()
+
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                timestep_emb,
+                post_patch_height,
+                post_patch_width,
+                encoder_attention_mask=encoder_attention_mask,
+            )
+        hidden_states = self.norm_out(
+            hidden_states, embedded_timestep, self.scale_shift_table
+        )
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_height, post_patch_width, p, p, self.out_channels
+        )
+        hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4)
+        hidden_states = hidden_states.reshape(
+            batch_size, self.out_channels, height, width
+        )
+
+        return hidden_states
+
+
+EntryClass = SanaTransformer2DModel

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

@@ -0,0 +1,906 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Callable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.configs.models.dits.sana_wm import SanaWMConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+
+# Re-exported for back-compat: callers import these names from this module path.
+from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import (  # noqa: F401
+    _CACHE_TYPE_CONCAT,
+    _CACHE_TYPE_STATE,
+    _INT32_SAFE_CONV_ELEMENTS,
+    _NUM_STREAM_CACHE_SLOTS,
+    _SLOT_CAM_K,
+    _SLOT_CAM_V,
+    _SLOT_FFN_TCONV,
+    _SLOT_K,
+    _SLOT_SHORTCONV,
+    _SLOT_TYPE_FLAG,
+    _SLOT_V,
+    BidirectionalGDNUCPESinglePathLiteLA,
+    CaptionEmbedder,
+    GLUMBConvTemp,
+    MultiHeadCrossAttention,
+    PatchEmbedMS3D,
+    T2IFinalLayer,
+    TimestepEmbedder,
+    WanRotaryPosEmbed,
+    _apply_block_diagonal,
+    _apply_complex_rope,
+    _apply_ray_projmat,
+    _apply_rotary_emb_bhnd,
+    _apply_rotary_emb_dn,
+    _bidirectional_short_conv,
+    _build_ucpe_apply_fns,
+    _compute_fov_from_focal,
+    _compute_frame_gates,
+    _ConvLayer,
+    _downscale_to_reference_rms,
+    _flip_and_shift,
+    _gdn_chunk_scan_forward,
+    _gdn_scan_bidirectional,
+    _gdn_scan_cached,
+    _gdn_scan_forward,
+    _gdn_scan_forward_stateful,
+    _invert_SE3,
+    _log_sana_wm_triton_cam_gdn_fallback,
+    _log_sana_wm_triton_gdn_fallback,
+    _RMSNorm,
+    _sana_wm_chunk_boundaries_for_attention,
+    _sana_wm_chunk_index_from_chunk_size,
+    _sana_wm_chunked_attention,
+    _sana_wm_normalize_chunk_index,
+    _sana_wm_padded_scale,
+    _sana_wm_sdpa,
+    _ShortConvolution,
+    _single_path_delta_chunk_scan_forward,
+    _single_path_delta_scan_bidirectional,
+    _single_path_delta_scan_cached,
+    _single_path_delta_scan_forward,
+    _single_path_delta_scan_forward_stateful,
+    _sinusoidal_timestep_embedding,
+    _slice_rope_for_cam,
+    _slice_rope_to_current_chunk,
+    _temporal_short_conv_cached,
+    _tensor_cache_key,
+    _UpstreamMlp,
+    compute_chunk_plucker,
+    process_camera_conditions_ucpe,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import (
+    parity_probe,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SanaWMBlock(nn.Module):
+    """One transformer block of SANA-WM."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        head_dim: int,
+        mlp_ratio: float,
+        t_kernel_size: int,
+        qk_norm: bool,
+        cross_norm: bool,
+        conv_kernel_size: int,
+        k_conv_only: bool,
+        softmax_main: bool,
+        use_chunk_plucker_post_attn: bool,
+        chunk_size: Optional[int] = None,
+        chunk_split_strategy: str = "uniform",
+        update_rule: str = "torch_chunk",
+        cam_update_rule: str = "torch_chunk",
+        chunk_gdn_chunk_size: int = 21,
+        use_chunked_softmax_attention: bool = False,
+        gdn_backend: str = "auto",
+    ) -> None:
+        super().__init__()
+        self.softmax_main = softmax_main
+        self.chunk_size = chunk_size
+        self.chunk_split_strategy = chunk_split_strategy
+
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.attn = BidirectionalGDNUCPESinglePathLiteLA(
+            in_dim=hidden_size,
+            heads=num_heads,
+            head_dim=head_dim,
+            qk_norm=qk_norm,
+            conv_kernel_size=conv_kernel_size,
+            k_conv_only=k_conv_only,
+            softmax_main=softmax_main,
+            update_rule=update_rule,
+            cam_update_rule=cam_update_rule,
+            chunk_gdn_chunk_size=chunk_gdn_chunk_size,
+            use_chunked_softmax_attention=use_chunked_softmax_attention,
+            gdn_backend=gdn_backend,
+        )
+
+        self.cross_attn = MultiHeadCrossAttention(
+            d_model=hidden_size,
+            num_heads=num_heads,
+            qk_norm=cross_norm,
+        )
+
+        self.mlp = GLUMBConvTemp(
+            in_features=hidden_size,
+            hidden_features=int(hidden_size * mlp_ratio),
+            t_kernel_size=t_kernel_size,
+        )
+
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(6, hidden_size) / hidden_size**0.5
+        )
+
+        if use_chunk_plucker_post_attn:
+            self.plucker_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+            nn.init.zeros_(self.plucker_proj.weight)
+            nn.init.zeros_(self.plucker_proj.bias)
+        else:
+            self.plucker_proj = None
+
+    @staticmethod
+    def _modulate(
+        x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
+    ) -> torch.Tensor:
+        return x * (1 + scale) + shift
+
+    @staticmethod
+    def _reshape_framewise_modulation(
+        x: torch.Tensor,
+        num_frames: int,
+    ) -> tuple[torch.Tensor, int]:
+        B, N, C = x.shape
+        tokens_per_frame = N // num_frames
+        return x.reshape(B, num_frames, tokens_per_frame, C), tokens_per_frame
+
+    def forward(
+        self,
+        x: torch.Tensor,  # (B, N, D)
+        y: torch.Tensor,  # (B, L, D) text embeds
+        t: torch.Tensor,  # (B, 6*D) AdaLN-single
+        HW: Tuple[int, int, int],
+        rotary_emb: Optional[torch.Tensor],
+        prope_fns: Optional[Tuple[Callable, Callable, Callable]],
+        plucker_emb: Optional[torch.Tensor],
+        mask: Optional[torch.Tensor],
+        chunk_size: Optional[int] = None,
+        chunk_split_strategy: Optional[str] = None,
+        chunk_index: Optional[List[int]] = None,
+    ) -> torch.Tensor:
+        B = x.shape[0]
+        if t.dim() == 2:
+            num_frames = None
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + t.reshape(B, 6, -1)
+            ).chunk(6, dim=1)
+        else:
+            num_frames = t.reshape(B, -1, 6, t.shape[-1] // 6).shape[1]
+            t = t.reshape(B, num_frames, 6, -1)
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None, None, :, :] + t
+            ).chunk(6, dim=2)
+
+        # Self-attention with UCPE camera branch
+        if num_frames is None:
+            x_in = self._modulate(self.norm1(x), shift_msa, scale_msa)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm1(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_msa, scale_msa).reshape_as(x)
+        attn_out = self.attn(
+            x_in,
+            HW=HW,
+            rotary_emb=rotary_emb,
+            prope_fns=prope_fns,
+            chunk_size=self.chunk_size if chunk_size is None else chunk_size,
+            chunk_split_strategy=(
+                self.chunk_split_strategy
+                if chunk_split_strategy is None
+                else chunk_split_strategy
+            ),
+            chunk_index=chunk_index,
+        )
+        if num_frames is None:
+            x = x + gate_msa * attn_out
+        else:
+            attn_out = attn_out.reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_msa * attn_out).reshape_as(x)
+
+        # Plücker post-attn injection (zero-init linear)
+        if self.plucker_proj is not None and plucker_emb is not None:
+            x = x + self.plucker_proj(plucker_emb)
+
+        # Cross-attention
+        x = x + self.cross_attn(x, y, mask=mask)
+
+        # FFN
+        if num_frames is None:
+            x_in = self._modulate(self.norm2(x), shift_mlp, scale_mlp)
+            x = x + gate_mlp * self.mlp(x_in, HW=HW)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm2(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_mlp, scale_mlp).reshape_as(x)
+            mlp_out = self.mlp(x_in, HW=HW).reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_mlp * mlp_out).reshape_as(x)
+        return x
+
+    def forward_long(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        t: torch.Tensor,
+        HW: Tuple[int, int, int],
+        rotary_emb: Optional[torch.Tensor],
+        prope_fns: Optional[Tuple[Callable, Callable, Callable]],
+        plucker_emb: Optional[torch.Tensor],
+        mask: Optional[torch.Tensor],
+        *,
+        kv_cache: list,
+        save_kv_cache: bool,
+    ) -> Tuple[torch.Tensor, list]:
+        """Streaming counterpart of ``forward``: threads the per-block 10-slot ``kv_cache`` through cached attention + FFN."""
+        B = x.shape[0]
+        if t.dim() == 2:
+            num_frames = None
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + t.reshape(B, 6, -1)
+            ).chunk(6, dim=1)
+        else:
+            num_frames = t.reshape(B, -1, 6, t.shape[-1] // 6).shape[1]
+            t = t.reshape(B, num_frames, 6, -1)
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None, None, :, :] + t
+            ).chunk(6, dim=2)
+
+        if num_frames is None:
+            x_in = self._modulate(self.norm1(x), shift_msa, scale_msa)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm1(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_msa, scale_msa).reshape_as(x)
+
+        attn_out, kv_cache = self.attn.forward_long(
+            x_in,
+            HW=HW,
+            rotary_emb=rotary_emb,
+            prope_fns=prope_fns,
+            kv_cache=kv_cache,
+            save_kv_cache=save_kv_cache,
+        )
+        if num_frames is None:
+            x = x + gate_msa * attn_out
+        else:
+            attn_out = attn_out.reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_msa * attn_out).reshape_as(x)
+
+        if self.plucker_proj is not None and plucker_emb is not None:
+            x = x + self.plucker_proj(plucker_emb)
+
+        x = x + self.cross_attn(x, y, mask=mask)
+
+        if num_frames is None:
+            x_in = self._modulate(self.norm2(x), shift_mlp, scale_mlp)
+        else:
+            x_norm, tokens_per_frame = self._reshape_framewise_modulation(
+                self.norm2(x), num_frames
+            )
+            x_in = self._modulate(x_norm, shift_mlp, scale_mlp).reshape_as(x)
+
+        # GLUMBConvTemp returns a tuple whenever the streaming path is active
+        # (ffn_tail set OR save requested); branch on tuple-ness, never on
+        # save_kv_cache alone (a read-only pass with a populated slot 9 still
+        # returns a tuple).
+        mlp_out = self.mlp(
+            x_in,
+            HW=HW,
+            ffn_tail=kv_cache[_SLOT_FFN_TCONV],
+            save_ffn_tail=save_kv_cache,
+        )
+        if isinstance(mlp_out, tuple):
+            mlp_out, ffn_tail = mlp_out
+            if save_kv_cache:
+                kv_cache[_SLOT_FFN_TCONV] = ffn_tail
+        if num_frames is None:
+            x = x + gate_mlp * mlp_out
+        else:
+            mlp_out = mlp_out.reshape(B, num_frames, tokens_per_frame, -1)
+            x = x + (gate_mlp * mlp_out).reshape_as(x)
+        return x, kv_cache
+
+
+class SanaWMTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """SANA-WM 2.6B TI2V world model.
+
+    Forward inputs:
+        hidden_states:           (B, C, T, H, W)        128-ch LTX-2 latent
+        encoder_hidden_states:   (B, L, 2304)           Gemma-2 embeddings
+        timestep:                (B,)
+        encoder_attention_mask:  (B, L) optional bool
+        camera_conditions:       (B, T, 20)             latent-frame raymap:
+                                                        16 c2w + (fx,fy,cx,cy)
+        chunk_plucker:           (B, 48, T, H, W)       optional, computed
+                                                        from camera_conditions
+                                                        if absent.
+
+    Returns: ``(B, C, T, H, W)`` predicted velocity / noise.
+    """
+
+    _fsdp_shard_conditions = SanaWMConfig()._fsdp_shard_conditions
+    _compile_conditions = SanaWMConfig()._compile_conditions
+    _supported_attention_backends = SanaWMConfig()._supported_attention_backends
+    param_names_mapping = SanaWMConfig().param_names_mapping
+    reverse_param_names_mapping = SanaWMConfig().reverse_param_names_mapping
+    lora_param_names_mapping: dict = {}
+
+    def __init__(self, config: SanaWMConfig, hf_config=None, **kwargs) -> None:
+        super().__init__(config, hf_config=hf_config or {}, **kwargs)
+        arch = config.arch_config
+
+        self.patch_size = (arch.patch_size_t, arch.patch_size, arch.patch_size)
+        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
+        self.hidden_size = self.inner_dim
+        self.num_attention_heads = arch.num_attention_heads
+        self.attention_head_dim = arch.attention_head_dim
+        self.out_channels = arch.out_channels
+        self.num_channels_latents = arch.num_channels_latents
+        self.vae_temporal_stride = arch.vae_temporal_stride
+        self.timestep_norm_scale_factor = getattr(
+            arch, "timestep_norm_scale_factor", 1.0
+        )
+
+        # --- Embedders ---
+        self.x_embedder = PatchEmbedMS3D(
+            self.patch_size,
+            arch.in_channels,
+            self.inner_dim,
+            bias=True,
+        )
+
+        self.t_embedder = TimestepEmbedder(self.inner_dim, frequency_embedding_size=256)
+        self.t_block = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True),
+        )
+
+        self.y_embedder = CaptionEmbedder(
+            in_channels=arch.caption_channels,
+            hidden_size=self.inner_dim,
+            token_num=arch.model_max_length,
+        )
+        self.y_norm = bool(getattr(arch, "y_norm", True))
+        self.attention_y_norm = _RMSNorm(
+            self.inner_dim,
+            scale_factor=getattr(arch, "y_norm_scale_factor", 1.0),
+            eps=getattr(arch, "y_norm_eps", 1e-5),
+        )
+
+        # 3-channel raymap embedder -- kept for state_dict compatibility but
+        # only invoked when ``use_chunk_plucker_post_attn`` is False.
+        # When ``True`` (the case for the released checkpoint) the absmap
+        # path is skipped entirely.
+        self.raymap_embedder = PatchEmbedMS3D(
+            self.patch_size,
+            3,
+            self.inner_dim,
+            bias=True,
+        )
+        # 48-channel plucker embedder (chunk-packed)
+        if arch.use_chunk_plucker_post_attn or arch.use_chunk_plucker_input:
+            self.plucker_embedder = PatchEmbedMS3D(
+                self.patch_size,
+                arch.chunk_plucker_channels,
+                self.inner_dim,
+                bias=True,
+            )
+            nn.init.zeros_(self.plucker_embedder.proj.weight)
+            nn.init.zeros_(self.plucker_embedder.proj.bias)
+        else:
+            self.plucker_embedder = None
+        self.use_chunk_plucker_post_attn = arch.use_chunk_plucker_post_attn
+        self.use_chunk_plucker_input = arch.use_chunk_plucker_input
+        self.chunk_size = getattr(arch, "chunk_size", None)
+        self.chunk_split_strategy = getattr(arch, "chunk_split_strategy", "uniform")
+
+        # --- RoPE ---
+        self.rope = WanRotaryPosEmbed(
+            attention_head_dim=arch.linear_head_dim,
+            patch_size=self.patch_size,
+            max_seq_len=1024,
+        )
+
+        # --- Transformer blocks ---
+        depth = arch.num_layers
+        self.softmax_every_n = arch.softmax_every_n
+        softmax_idx = set(
+            i
+            for i in range(depth)
+            if arch.softmax_every_n > 0 and (i + 1) % arch.softmax_every_n == 0
+        )
+        self.softmax_block_indices = tuple(sorted(softmax_idx))
+
+        self.blocks = nn.ModuleList(
+            [
+                SanaWMBlock(
+                    hidden_size=self.inner_dim,
+                    num_heads=arch.num_attention_heads,
+                    head_dim=arch.linear_head_dim,
+                    mlp_ratio=arch.mlp_ratio,
+                    t_kernel_size=arch.t_kernel_size,
+                    qk_norm=arch.qk_norm,
+                    cross_norm=arch.cross_norm,
+                    conv_kernel_size=arch.conv_kernel_size,
+                    k_conv_only=arch.k_conv_only,
+                    softmax_main=(i in softmax_idx),
+                    use_chunk_plucker_post_attn=(
+                        arch.use_chunk_plucker_post_attn
+                        and (
+                            arch.chunk_plucker_post_attn_blocks < 0
+                            or i < arch.chunk_plucker_post_attn_blocks
+                        )
+                    ),
+                    chunk_size=self.chunk_size,
+                    chunk_split_strategy=self.chunk_split_strategy,
+                    update_rule=getattr(arch, "update_rule", "torch_chunk"),
+                    cam_update_rule=getattr(arch, "cam_update_rule", "torch_chunk"),
+                    chunk_gdn_chunk_size=getattr(arch, "chunk_gdn_chunk_size", 21),
+                    use_chunked_softmax_attention=getattr(
+                        arch, "use_chunked_softmax_attention", False
+                    ),
+                    gdn_backend=getattr(arch, "gdn_backend", "auto"),
+                )
+                for i in range(depth)
+            ]
+        )
+
+        self.final_layer = T2IFinalLayer(
+            self.inner_dim, self.patch_size, self.out_channels
+        )
+
+        # Cache RoPE freqs per shape -- avoids recomputation across denoising
+        # steps with constant latent shapes.
+        self._freqs_cache: dict = {}
+        self._ucpe_apply_fns_cache: Optional[
+            Tuple[Tuple, torch.Tensor, Tuple[Callable, Callable, Callable]]
+        ] = None
+        self._plucker_emb_cache: Optional[Tuple[Tuple, torch.Tensor, torch.Tensor]] = (
+            None
+        )
+
+        # FSDP shard targets
+        self.layer_names = ["blocks"]
+
+    def post_load_weights(self) -> None:
+        # FSDP loader initializes the model on meta and only materializes
+        # tensors that appear in the checkpoint. WanRotaryPosEmbed._freqs is a
+        # derived, non-persistent constant, so recompute it deterministically.
+        for module in self.modules():
+            if isinstance(module, WanRotaryPosEmbed):
+                if module._freqs.is_meta:
+                    module._init_freqs_buffer()
+
+    # ------------------------------------------------------------------ #
+    # Forward
+    # ------------------------------------------------------------------ #
+
+    def _get_freqs(self, T: int, H: int, W: int, device: torch.device) -> torch.Tensor:
+        key = (T, H, W, str(device))
+        if key not in self._freqs_cache:
+            self._freqs_cache[key] = self.rope((T, H, W), device)
+        return self._freqs_cache[key]
+
+    def _get_freqs_window(
+        self,
+        start: int,
+        end: int,
+        H: int,
+        W: int,
+        device: torch.device,
+        frame_index: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """RoPE freqs for a streaming chunk at GLOBAL frame positions.
+
+        ``frame_index`` (per-token global positions) overrides ``(start, end)``;
+        the count branch is cached, the tensor branch is computed fresh.
+        """
+        if frame_index is not None:
+            return self.rope((end - start, H, W), device, frame_index=frame_index)
+        key = ("win", int(start), int(end), H, W, str(device))
+        if key not in self._freqs_cache:
+            self._freqs_cache[key] = self.rope(((int(start), int(end)), H, W), device)
+        return self._freqs_cache[key]
+
+    def _get_ucpe_apply_fns(
+        self,
+        camera_conditions: torch.Tensor,
+        *,
+        HW: Tuple[int, int, int],
+        freqs: torch.Tensor,
+    ) -> Tuple[Callable, Callable, Callable]:
+        head_dim = self.attention_head_dim
+        if torch.is_grad_enabled():
+            raymats = process_camera_conditions_ucpe(
+                camera_conditions,
+                HW=HW,
+                patch_size=self.patch_size,
+            )
+            raymats_flat = raymats.reshape(camera_conditions.shape[0], -1, 4, 4)
+            return _build_ucpe_apply_fns(head_dim, raymats_flat, freqs)
+
+        key = (
+            "ucpe",
+            HW,
+            self.patch_size,
+            head_dim,
+            _tensor_cache_key(camera_conditions),
+            _tensor_cache_key(freqs),
+        )
+        cached = self._ucpe_apply_fns_cache
+        if cached is not None and cached[0] == key:
+            return cached[2]
+
+        raymats = process_camera_conditions_ucpe(
+            camera_conditions,
+            HW=HW,
+            patch_size=self.patch_size,
+        )
+        raymats_flat = raymats.reshape(camera_conditions.shape[0], -1, 4, 4)
+        prope_fns = _build_ucpe_apply_fns(head_dim, raymats_flat, freqs)
+        self._ucpe_apply_fns_cache = (key, camera_conditions, prope_fns)
+        return prope_fns
+
+    def _get_plucker_emb(
+        self,
+        chunk_plucker: torch.Tensor,
+        *,
+        latent_token_count: int,
+    ) -> torch.Tensor:
+        if self.plucker_embedder is None:
+            raise ValueError("SANA-WM plucker_embedder is not initialized.")
+
+        weight = self.plucker_embedder.proj.weight
+        bias = self.plucker_embedder.proj.bias
+        key = (
+            "plucker_emb",
+            latent_token_count,
+            self.patch_size,
+            _tensor_cache_key(chunk_plucker),
+            _tensor_cache_key(weight),
+            None if bias is None else _tensor_cache_key(bias),
+        )
+        if not torch.is_grad_enabled():
+            cached = self._plucker_emb_cache
+            if cached is not None and cached[0] == key:
+                return cached[2]
+
+        plucker_emb = self.plucker_embedder(chunk_plucker.to(weight.dtype))
+        if plucker_emb.shape[1] != latent_token_count:
+            raise ValueError(
+                f"plucker_emb token count {plucker_emb.shape[1]} != "
+                f"latent token count {latent_token_count}; "
+                "expected chunk_plucker shape (B, 48, T, H, W)."
+            )
+
+        if not torch.is_grad_enabled():
+            self._plucker_emb_cache = (key, chunk_plucker, plucker_emb)
+        return plucker_emb
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        camera_conditions: Optional[torch.Tensor] = None,
+        chunk_plucker: Optional[torch.Tensor] = None,
+        guidance: Optional[torch.Tensor] = None,  # kept for compat
+        **kwargs,
+    ) -> torch.Tensor:
+        if encoder_hidden_states is None:
+            raise ValueError("SANA-WM forward requires encoder_hidden_states.")
+        if timestep is None:
+            raise ValueError("SANA-WM forward requires timestep.")
+
+        B, C, T_raw, H_raw, W_raw = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        T = T_raw // p_t
+        H = H_raw // p_h
+        W = W_raw // p_w
+        chunk_size = kwargs.get("chunk_size", self.chunk_size)
+        chunk_split_strategy = kwargs.get(
+            "chunk_split_strategy", self.chunk_split_strategy
+        )
+        chunk_index = kwargs.get("chunk_index", None)
+
+        # Patch embed: (B, C, T, H, W) -> (B, T*H*W, D)
+        x = self.x_embedder(hidden_states.to(dtype=self.x_embedder.proj.weight.dtype))
+
+        # Timestep AdaLN-single. SANA-WM's LTX sampler passes per-frame
+        # timesteps shaped (B, 1, T) so the clean first-frame condition can stay
+        # at timestep 0 while remaining latent frames denoise. Keep the scalar
+        # path for generic scheduler compatibility.
+        if self.timestep_norm_scale_factor != 1.0:
+            timestep_for_embed = (
+                timestep.float() / self.timestep_norm_scale_factor
+            ).to(torch.float32)
+        else:
+            timestep_for_embed = timestep.long().to(torch.float32)
+
+        if timestep_for_embed.dim() == 1:
+            t_emb = self.t_embedder(timestep_for_embed)  # (B, D)
+            t6 = self.t_block(t_emb)  # (B, 6D)
+        else:
+            timestep_shape = tuple(timestep_for_embed.shape)
+            t_flat = self.t_embedder(timestep_for_embed.flatten())
+            t6_flat = self.t_block(t_flat)
+            t_emb = t_flat.unflatten(0, timestep_shape)
+            t6 = t6_flat.unflatten(0, timestep_shape)
+
+        if isinstance(encoder_attention_mask, (list, tuple)):
+            encoder_attention_mask = encoder_attention_mask[0]
+        y = encoder_hidden_states
+        if y.dim() == 3:
+            y = y.unsqueeze(1)
+        y = self.y_embedder(y).squeeze(1)  # (B, L, D)
+        if y.shape[0] != B:
+            y = y.expand(B, -1, -1).contiguous()
+        if self.y_norm:
+            y = self.attention_y_norm(y)
+        if encoder_attention_mask is not None and encoder_attention_mask.shape[0] != B:
+            encoder_attention_mask = encoder_attention_mask.expand(B, -1).contiguous()
+
+        freqs = self._get_freqs(T, H, W, x.device)
+
+        # Camera conditioning: UCPE prope_fns + Plücker
+        prope_fns = None
+        if camera_conditions is not None:
+            if camera_conditions.shape[1] != T:
+                raise ValueError(
+                    "SANA-WM camera_conditions must be sampled at latent "
+                    f"frames: got {camera_conditions.shape[1]} frames, "
+                    f"expected T={T}."
+                )
+            prope_fns = self._get_ucpe_apply_fns(
+                camera_conditions,
+                HW=(T, H, W),
+                freqs=freqs,
+            )
+
+        # Plücker post-attn embedding (shared across all blocks)
+        plucker_emb = None
+        needs_plucker_emb = (
+            chunk_plucker is not None
+            and self.plucker_embedder is not None
+            and (self.use_chunk_plucker_post_attn or self.use_chunk_plucker_input)
+        )
+        if needs_plucker_emb:
+            plucker_emb = self._get_plucker_emb(
+                chunk_plucker,
+                latent_token_count=x.shape[1],
+            )  # (B, T*H*W, D)
+
+        if self.use_chunk_plucker_input and plucker_emb is not None:
+            x = x + plucker_emb
+
+        if not self.use_chunk_plucker_post_attn:
+            plucker_emb = None
+
+        # --- 6. Transformer blocks ---
+        HW = (T, H, W)
+        for block in self.blocks:
+            x = block(
+                x,
+                y=y,
+                t=t6,
+                HW=HW,
+                rotary_emb=freqs,
+                prope_fns=prope_fns,
+                plucker_emb=plucker_emb,
+                mask=encoder_attention_mask,
+                chunk_size=chunk_size,
+                chunk_split_strategy=chunk_split_strategy,
+                chunk_index=chunk_index,
+            )
+
+        x = self.final_layer(x, t_emb)  # (B, N, p_t*p_h*p_w*C_out)
+
+        # Un-patch
+        x = x.reshape(B, T, H, W, p_t, p_h, p_w, self.out_channels)
+        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).contiguous()
+        x = x.reshape(B, self.out_channels, T * p_t, H * p_h, W * p_w)
+        return x
+
+    def forward_long(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        camera_conditions: Optional[torch.Tensor] = None,
+        chunk_plucker: Optional[torch.Tensor] = None,
+        *,
+        kv_cache: Optional[list] = None,
+        save_kv_cache: bool = True,
+        start_f: Optional[int] = None,
+        end_f: Optional[int] = None,
+        frame_index: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, list]:
+        """Streaming autoregressive forward over a chunk of latent frames.
+
+        RoPE / camera / plücker are windowed to the chunk's GLOBAL frame range
+        ``[start_f, end_f)``; a per-block 10-slot ``kv_cache`` carries recurrent
+        state / concat-windows across chunks. Returns ``(out, new_cache)``.
+        """
+        if encoder_hidden_states is None:
+            raise ValueError("SANA-WM forward_long requires encoder_hidden_states.")
+        if timestep is None:
+            raise ValueError("SANA-WM forward_long requires timestep.")
+
+        if kv_cache is None:
+            kv_cache = [[None] * _NUM_STREAM_CACHE_SLOTS for _ in self.blocks]
+
+        B, C, T_raw, H_raw, W_raw = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        T = T_raw // p_t
+        H = H_raw // p_h
+        W = W_raw // p_w
+        start = 0 if start_f is None else int(start_f)
+        end = start + T if end_f is None else int(end_f)
+
+        x = self.x_embedder(hidden_states.to(dtype=self.x_embedder.proj.weight.dtype))
+
+        # Timestep AdaLN-single: force the framewise (B, 1, T) path so blocks
+        # always apply per-frame modulation.
+        if timestep.dim() == 1:
+            timestep = timestep[:, None, None].expand(-1, 1, T)
+        elif timestep.dim() == 2:
+            timestep = timestep[:, None, :]
+        if self.timestep_norm_scale_factor != 1.0:
+            timestep_for_embed = (
+                timestep.float() / self.timestep_norm_scale_factor
+            ).to(torch.float32)
+        else:
+            timestep_for_embed = timestep.long().to(torch.float32)
+        timestep_shape = tuple(timestep_for_embed.shape)
+        t_flat = self.t_embedder(timestep_for_embed.flatten())
+        t6_flat = self.t_block(t_flat)
+        t_emb = t_flat.unflatten(0, timestep_shape)
+        t6 = t6_flat.unflatten(0, timestep_shape)
+
+        if isinstance(encoder_attention_mask, (list, tuple)):
+            encoder_attention_mask = encoder_attention_mask[0]
+        y = encoder_hidden_states
+        if y.dim() == 3:
+            y = y.unsqueeze(1)
+        y = self.y_embedder(y).squeeze(1)
+        if y.shape[0] != B:
+            y = y.expand(B, -1, -1).contiguous()
+        if self.y_norm:
+            y = self.attention_y_norm(y)
+        if encoder_attention_mask is not None and encoder_attention_mask.shape[0] != B:
+            encoder_attention_mask = encoder_attention_mask.expand(B, -1).contiguous()
+
+        # RoPE windowed to global frame positions [start, end)
+        freqs = self._get_freqs_window(
+            start, end, H, W, x.device, frame_index=frame_index
+        )
+
+        # Camera conditioning: slice to the chunk, co-windowed w/ freqs
+        prope_fns = None
+        if camera_conditions is not None:
+            if camera_conditions.shape[1] != T:
+                # .contiguous(): canonical layout regardless of how the caller
+                # built the full-length tensor, so batch and realtime windows are
+                # kernel-level identical (slice offset/stride changes the reduction
+                # order otherwise — measured 1e-7 seeds amplifying to %-level drift
+                # through the bf16 block stack).
+                camera_conditions = camera_conditions[:, start:end].contiguous()
+            if camera_conditions.shape[0] != B:
+                camera_conditions = camera_conditions.repeat(
+                    B // camera_conditions.shape[0], 1, 1
+                )
+            prope_fns = self._get_ucpe_apply_fns(
+                camera_conditions, HW=(T, H, W), freqs=freqs
+            )
+
+        # Plücker post-attn / input embedding, sliced to the chunk.
+        if chunk_plucker is not None and chunk_plucker.shape[2] != T:
+            chunk_plucker = chunk_plucker[
+                :, :, start:end
+            ].contiguous()  # see camera note
+        if chunk_plucker is not None and chunk_plucker.shape[0] != B:
+            chunk_plucker = chunk_plucker.repeat(
+                B // chunk_plucker.shape[0], 1, 1, 1, 1
+            )
+        plucker_emb = None
+        needs_plucker_emb = (
+            chunk_plucker is not None
+            and self.plucker_embedder is not None
+            and (self.use_chunk_plucker_post_attn or self.use_chunk_plucker_input)
+        )
+        if needs_plucker_emb:
+            plucker_emb = self._get_plucker_emb(
+                chunk_plucker, latent_token_count=x.shape[1]
+            )
+        if self.use_chunk_plucker_input and plucker_emb is not None:
+            x = x + plucker_emb
+        if not self.use_chunk_plucker_post_attn:
+            plucker_emb = None
+
+        # parity harness (env-gated, no-op in prod): on the FIRST sink-path call
+        # (frame_index not None), checksum the pre-block tensors and x after every
+        # block to localize where the two execution paths first diverge.
+        _probe_path = os.environ.get(parity_probe.ENV_BLOCK_PROBE)
+        _probe = None
+        if (
+            _probe_path
+            and frame_index is not None
+            and not getattr(self, "_block_probe_done", False)
+        ):
+            _ck = parity_probe.checksum
+            _probe = {
+                "x_embed": _ck(x),
+                "t6": _ck(t6),
+                "y": _ck(y),
+                "freqs": (
+                    (
+                        tuple(freqs.shape),
+                        float(freqs.real.detach().double().sum().item()),
+                        float(freqs.imag.detach().double().sum().item()),
+                    )
+                    if freqs is not None
+                    else None
+                ),
+                "plucker_emb": _ck(plucker_emb),
+                "frame_index": frame_index.tolist(),
+            }
+
+        HW = (T, H, W)
+        new_cache = []
+        for i, block in enumerate(self.blocks):
+            x, block_cache = block.forward_long(
+                x,
+                y,
+                t6,
+                HW,
+                freqs,
+                prope_fns,
+                plucker_emb,
+                encoder_attention_mask,
+                kv_cache=kv_cache[i],
+                save_kv_cache=save_kv_cache,
+            )
+            new_cache.append(block_cache)
+            if _probe is not None:
+                _probe[f"x_after_block_{i:02d}"] = parity_probe.checksum(x)
+        if _probe is not None:
+            torch.save(_probe, _probe_path)
+            self._block_probe_done = True
+
+        x = self.final_layer(x, t_emb)
+        x = x.reshape(B, T, H, W, p_t, p_h, p_w, self.out_channels)
+        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).contiguous()
+        x = x.reshape(B, self.out_channels, T * p_t, H * p_h, W * p_w)
+        return x, new_cache
+
+
+EntryClass = SanaWMTransformer3DModel

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

@@ -0,0 +1,426 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.configs.models.dits.sana_wm_refiner import (
+    SanaWMRefinerArchConfig,
+    SanaWMRefinerConfig,
+)
+from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.models.dits.ltx_2 import (
+    LTX2AdaLayerNormSingle,
+    LTX2Attention,
+    LTX2AudioVideoRotaryPosEmbed,
+    LTX2FeedForward,
+    LTX2TextProjection,
+)
+
+
+def pack_latents(
+    latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1
+) -> torch.Tensor:
+    """Pack a 5D latent (B, C, T, H, W) into a 3D token sequence (B, L, in_dim)."""
+    B, _, T, H, W = latents.shape
+    pT = T // patch_size_t
+    pH = H // patch_size
+    pW = W // patch_size
+    latents = latents.reshape(B, -1, pT, patch_size_t, pH, patch_size, pW, patch_size)
+    return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+
+
+def unpack_latents(
+    tokens: torch.Tensor,
+    num_frames: int,
+    height: int,
+    width: int,
+    patch_size: int = 1,
+    patch_size_t: int = 1,
+) -> torch.Tensor:
+    """Inverse of `pack_latents`: (B, L, out_dim) -> (B, C, T, H, W)."""
+    B = tokens.size(0)
+    tokens = tokens.reshape(
+        B,
+        num_frames // patch_size_t,
+        height // patch_size,
+        width // patch_size,
+        -1,
+        patch_size_t,
+        patch_size,
+        patch_size,
+    )
+    return (
+        tokens.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
+    )
+
+
+def _slice_rope(
+    rope: tuple[torch.Tensor, torch.Tensor], start: int, end: Optional[int] = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Slice along token axis for either interleaved (rank-3) or split (rank-4)."""
+    cos, sin = rope
+    end_ = end if end is not None else cos.shape[-2 if cos.ndim == 4 else 1]
+    if cos.ndim == 3:
+        return cos[:, start:end_], sin[:, start:end_]
+    if cos.ndim == 4:
+        return cos[:, :, start:end_, :], sin[:, :, start:end_, :]
+    raise ValueError(f"Unexpected RoPE rank: {cos.ndim}")
+
+
+def _streaming_self_attention(
+    attn: LTX2Attention,
+    hidden_states: torch.Tensor,
+    video_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    n_context_tokens: int,
+) -> torch.Tensor:
+    """Streaming SLA: context attends to context only, current attends to context+current.
+
+    Mirrors NVlabs `inference_sana_wm.py::_streaming_self_attention`.
+    """
+    seq_len = hidden_states.shape[1]
+    if n_context_tokens <= 0 or n_context_tokens >= seq_len:
+        return attn(hidden_states, context=None, pe=video_rotary_emb)
+
+    ctx_rope = _slice_rope(video_rotary_emb, 0, n_context_tokens)
+    out_ctx = attn(
+        hidden_states[:, :n_context_tokens],
+        context=None,
+        pe=ctx_rope,
+    )
+
+    cur_rope = _slice_rope(video_rotary_emb, n_context_tokens, seq_len)
+    out_cur = attn(
+        hidden_states[:, n_context_tokens:],
+        context=hidden_states,
+        pe=cur_rope,
+        k_pe=video_rotary_emb,
+    )
+    return torch.cat([out_ctx, out_cur], dim=1)
+
+
+class SanaWMRefinerBlock(nn.Module):
+    """Video-only LTX-2 transformer block.
+
+    Diffusers-compatible layout: `norm1 -> attn1 (self) -> norm2 -> attn2 (cross) -> norm3 -> ff`,
+    each modulated via per-block `scale_shift_table` + token-wise `temb`.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        qk_norm: bool = True,
+        norm_eps: float = 1e-6,
+        apply_gated_attention: bool = False,
+        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__()
+        self.dim = int(dim)
+
+        self.norm1 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False)
+        self.attn1 = LTX2Attention(
+            query_dim=self.dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            norm_eps=norm_eps,
+            qk_norm=qk_norm,
+            apply_gated_attention=apply_gated_attention,
+            prefix=f"{prefix}.attn1",
+            quant_config=quant_config,
+        )
+
+        self.norm2 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False)
+        self.attn2 = LTX2Attention(
+            query_dim=self.dim,
+            context_dim=cross_attention_dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            norm_eps=norm_eps,
+            qk_norm=qk_norm,
+            use_local_attention=True,
+            apply_gated_attention=apply_gated_attention,
+            prefix=f"{prefix}.attn2",
+            quant_config=quant_config,
+        )
+
+        self.norm3 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False)
+        self.ff = LTX2FeedForward(self.dim, dim_out=self.dim, quant_config=quant_config)
+
+        self.scale_shift_table = nn.Parameter(torch.randn(6, self.dim) / self.dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        video_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        n_context_tokens: int = 0,
+    ) -> torch.Tensor:
+        B = hidden_states.size(0)
+        T = temb.size(1)
+        D = self.dim
+        ada = self.scale_shift_table[None, None].to(
+            device=temb.device, dtype=temb.dtype
+        ) + temb.reshape(B, T, 6, D)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada.unbind(
+            dim=2
+        )
+
+        normed = self.norm1(hidden_states) * (1 + scale_msa) + shift_msa
+        attn_out = _streaming_self_attention(
+            self.attn1,
+            normed,
+            video_rotary_emb,
+            n_context_tokens=n_context_tokens,
+        )
+        hidden_states = hidden_states + attn_out * gate_msa
+
+        normed = self.norm2(hidden_states)
+        ca_out = self.attn2(
+            normed,
+            context=encoder_hidden_states,
+            mask=encoder_attention_mask,
+            pe=None,
+        )
+        hidden_states = hidden_states + ca_out
+
+        normed = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp
+        hidden_states = hidden_states + self.ff(normed) * gate_mlp
+        return hidden_states
+
+
+class SanaWMLTX2VideoRefiner(CachableDiT, LayerwiseOffloadableModuleMixin):
+    """SANA-WM stage-2 LTX-2 video-only refiner.
+
+    Loads Diffusers-format refiner weights from `<model_path>/refiner/transformer/`.
+    Audio params present in the checkpoint are silently dropped by the loader's
+    `strict=False` state_dict load.
+    """
+
+    _fsdp_shard_conditions = SanaWMRefinerArchConfig()._fsdp_shard_conditions
+    _compile_conditions = SanaWMRefinerArchConfig()._compile_conditions
+    _supported_attention_backends = (
+        SanaWMRefinerArchConfig()._supported_attention_backends
+    )
+    param_names_mapping = SanaWMRefinerArchConfig().param_names_mapping
+    reverse_param_names_mapping: dict = {}
+    lora_param_names_mapping: dict = {}
+
+    def __init__(
+        self,
+        config: SanaWMRefinerConfig,
+        hf_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__(config, hf_config=hf_config)
+        arch = config.arch_config
+
+        self.in_channels = int(arch.in_channels)
+        self.out_channels = int(arch.out_channels)
+        self.patch_size = int(arch.patch_size)
+        self.patch_size_t = int(arch.patch_size_t)
+        self.hidden_size = int(arch.hidden_size)
+        self.num_attention_heads = int(arch.num_attention_heads)
+        self.num_channels_latents = int(arch.num_channels_latents)
+        self.attention_head_dim = int(arch.attention_head_dim)
+        self.timestep_scale_multiplier = float(arch.timestep_scale_multiplier)
+        self.rope_type = str(arch.rope_type)
+
+        in_dim = (
+            self.in_channels * self.patch_size_t * self.patch_size * self.patch_size
+        )
+        out_dim = (
+            self.out_channels * self.patch_size_t * self.patch_size * self.patch_size
+        )
+
+        self.proj_in = ColumnParallelLinear(
+            in_dim,
+            self.hidden_size,
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
+        )
+
+        self.time_embed = LTX2AdaLayerNormSingle(
+            self.hidden_size, embedding_coefficient=6
+        )
+        self.caption_projection = LTX2TextProjection(
+            in_features=int(arch.caption_channels),
+            hidden_size=self.hidden_size,
+            out_features=self.hidden_size,
+            act_fn="gelu_tanh",
+        )
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                SanaWMRefinerBlock(
+                    dim=self.hidden_size,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                    cross_attention_dim=int(arch.cross_attention_dim),
+                    qk_norm=bool(arch.qk_norm),
+                    norm_eps=float(arch.norm_eps),
+                    apply_gated_attention=bool(arch.apply_gated_attention),
+                    prefix=f"transformer_blocks.{i}",
+                    quant_config=quant_config,
+                )
+                for i in range(int(arch.num_layers))
+            ]
+        )
+
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, self.hidden_size) / self.hidden_size**0.5
+        )
+        self.norm_out = nn.LayerNorm(
+            self.hidden_size, eps=float(arch.norm_eps), elementwise_affine=False
+        )
+        self.proj_out = ColumnParallelLinear(
+            self.hidden_size,
+            out_dim,
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
+        )
+
+        # LTX2AudioVideoRotaryPosEmbed expects `dim` to be the *total* hidden
+        # size (num_heads * head_dim), not the per-head dim. It internally
+        # reshapes cos/sin to (B, T, num_heads, head_dim/2). Passing
+        # `attention_head_dim` here would size the RoPE to head_dim/num_heads
+        # and produce a (1, num_heads, L, 2) cos/sin that won't match
+        # LTX2Attention's q/k. See LTX2Transformer3DAVModel.__init__ in
+        # ltx_2.py for the canonical convention (`dim=self.hidden_size`).
+        self.rope = LTX2AudioVideoRotaryPosEmbed(
+            dim=self.hidden_size,
+            patch_size=self.patch_size,
+            patch_size_t=self.patch_size_t,
+            base_num_frames=int(arch.base_num_frames),
+            base_height=int(arch.base_height),
+            base_width=int(arch.base_width),
+            sampling_rate=int(arch.sampling_rate),
+            hop_length=int(arch.hop_length),
+            scale_factors=tuple(arch.scale_factors),
+            causal_offset=int(arch.causal_offset),
+            modality="video",
+            rope_type=self.rope_type,
+            num_attention_heads=self.num_attention_heads,
+        )
+
+        self.layer_names = ["transformer_blocks"]
+
+    def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor:
+        return timestep * self.timestep_scale_multiplier
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states_image=None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        fps: float = 24.0,
+        n_context_tokens: int = 0,
+        guidance=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        # Accept either packed (B, L, in_dim) or raw 5D (B, C, T, H, W).
+        if hidden_states.dim() == 5:
+            B_, _, T_, H_, W_ = hidden_states.shape
+            if num_frames is None:
+                num_frames = T_
+            if height is None:
+                height = H_
+            if width is None:
+                width = W_
+            hidden_states = pack_latents(
+                hidden_states,
+                patch_size=self.patch_size,
+                patch_size_t=self.patch_size_t,
+            )
+            packed_input = True
+        else:
+            if num_frames is None or height is None or width is None:
+                raise ValueError(
+                    "num_frames/height/width are required when hidden_states is pre-packed."
+                )
+            packed_input = False
+
+        B = hidden_states.size(0)
+
+        video_coords = self.rope.prepare_video_coords(
+            batch_size=B,
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            device=hidden_states.device,
+            fps=fps,
+        )
+        video_rotary_emb = self.rope(
+            video_coords,
+            device=hidden_states.device,
+            out_dtype=hidden_states.dtype,
+        )
+
+        hidden_states, _ = self.proj_in(hidden_states)
+
+        scaled_t = self._scale_timestep_for_adaln(timestep)
+        temb, embedded_timestep = self.time_embed(
+            scaled_t.flatten(), hidden_dtype=hidden_states.dtype
+        )
+        if timestep.dim() >= 2:
+            temb = temb.view(B, -1, temb.size(-1))
+            embedded_timestep = embedded_timestep.view(
+                B, -1, embedded_timestep.size(-1)
+            )
+        else:
+            temb = temb.view(B, 1, temb.size(-1))
+            embedded_timestep = embedded_timestep.view(B, 1, embedded_timestep.size(-1))
+
+        encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states.view(B, -1, self.hidden_size)
+
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                video_rotary_emb=video_rotary_emb,
+                encoder_attention_mask=encoder_attention_mask,
+                n_context_tokens=n_context_tokens,
+            )
+
+        scale_shift_values = self.scale_shift_table[None, None].to(
+            device=hidden_states.device, dtype=hidden_states.dtype
+        ) + embedded_timestep[:, :, None].to(hidden_states.dtype)
+        shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+        hidden_states = self.norm_out(hidden_states) * (1 + scale) + shift
+        hidden_states, _ = self.proj_out(hidden_states)
+
+        if packed_input:
+            return hidden_states
+        return unpack_latents(
+            hidden_states,
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            patch_size=self.patch_size,
+            patch_size_t=self.patch_size_t,
+        )
+
+
+EntryClass = SanaWMLTX2VideoRefiner

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

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+"""StableDiffusion3 Transformer model implementation.
+
+NOTE: This initial implementation uses diffusers' JointTransformerBlock directly.
+A native SGLang attention implementation is needed for FlashAttention, TP/SP,
+quantization, and LoRA support.
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import JointTransformerBlock
+from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
+from diffusers.models.normalization import AdaLayerNormContinuous
+
+from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import (
+    StableDiffusion3TransformerConfig,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SD3Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["JointTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+    layer_names = ["transformer_blocks"]
+
+    def __init__(
+        self,
+        config: StableDiffusion3TransformerConfig,
+        hf_config: dict[str, Any] | None = None,
+        quant_config=None,
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+        self.config = config
+        arch_config = config.arch_config
+        sample_size = arch_config.sample_size
+        patch_size = arch_config.patch_size
+        in_channels = arch_config.in_channels
+        num_layers = arch_config.num_layers
+        attention_head_dim = arch_config.attention_head_dim
+        num_attention_heads = arch_config.num_attention_heads
+        joint_attention_dim = arch_config.joint_attention_dim
+        caption_projection_dim = arch_config.caption_projection_dim
+        pooled_projection_dim = arch_config.pooled_projection_dim
+        out_channels = arch_config.out_channels
+        pos_embed_max_size = arch_config.pos_embed_max_size
+        dual_attention_layers = arch_config.dual_attention_layers
+        qk_norm = arch_config.qk_norm
+
+        self.out_channels = out_channels if out_channels is not None else in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.patch_size = patch_size
+
+        self.pos_embed = PatchEmbed(
+            height=sample_size,
+            width=sample_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            embed_dim=self.inner_dim,
+            pos_embed_max_size=pos_embed_max_size,
+        )
+        self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+            embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+        )
+        self.context_embedder = nn.Linear(joint_attention_dim, caption_projection_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                JointTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    context_pre_only=i == num_layers - 1,
+                    qk_norm=qk_norm,
+                    use_dual_attention=i in dual_attention_layers,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
+        )
+        self.proj_out = nn.Linear(
+            self.inner_dim, patch_size * patch_size * self.out_channels, bias=True
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        pooled_projections: torch.Tensor | None = None,
+        timestep: torch.LongTensor | None = None,
+        block_controlnet_hidden_states: list | None = None,
+        guidance: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        skip_layers: list[int] | None = None,
+    ) -> torch.Tensor:
+        if encoder_hidden_states is None:
+            raise ValueError("encoder_hidden_states must be provided.")
+        if pooled_projections is None:
+            raise ValueError("pooled_projections must be provided.")
+
+        encoder_embeddings = encoder_hidden_states
+
+        height, width = hidden_states.shape[-2:]
+
+        hidden_states = self.pos_embed(hidden_states)
+        temb = self.time_text_embed(timestep, pooled_projections)
+        encoder_embeddings = self.context_embedder(encoder_embeddings)
+
+        skip_layer_set = set(skip_layers) if skip_layers else set()
+
+        if block_controlnet_hidden_states is not None:
+            interval_control = len(self.transformer_blocks) / len(
+                block_controlnet_hidden_states
+            )
+        else:
+            interval_control = 0
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if index_block not in skip_layer_set:
+                encoder_embeddings, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_embeddings,
+                    temb=temb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if (
+                block_controlnet_hidden_states is not None
+                and block.context_pre_only is False
+            ):
+                hidden_states = (
+                    hidden_states
+                    + block_controlnet_hidden_states[
+                        int(index_block / interval_control)
+                    ]
+                )
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # unpatchify
+        patch_size = self.patch_size
+        height = height // patch_size
+        width = width // patch_size
+
+        hidden_states = hidden_states.reshape(
+            shape=(
+                hidden_states.shape[0],
+                height,
+                width,
+                patch_size,
+                patch_size,
+                self.out_channels,
+            )
+        )
+        hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4)
+        output = hidden_states.reshape(
+            shape=(
+                hidden_states.shape[0],
+                self.out_channels,
+                height * patch_size,
+                width * patch_size,
+            )
+        )
+
+        return output
+
+
+# Entry class for registry
+EntryClass = SD3Transformer2DModel

python/sglang/multimodal_gen/runtime/models/encoders/base.py

@@ -0,0 +1,170 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from abc import ABC, abstractmethod
+from dataclasses import field
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.encoders import (
+    BaseEncoderOutput,
+    EncoderConfig,
+    ImageEncoderConfig,
+    TextEncoderConfig,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+    get_sp_group,
+    get_tp_group,
+    get_world_group,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+def get_folding_tp_group(config: EncoderConfig):
+    """Group an encoder should tensor-parallel over.
+
+    ``config.parallel_folding_mode`` is set by ServerArgs.adjust_pipeline_config
+    when the encoder is folded over a larger group than its own TP (the idle DiT
+    replica during the encoding stage); when it is None the encoder uses the
+    default TP group. Shared by every text/image encoder so the choice lives in
+    one place.
+    """
+    mode = config.parallel_folding_mode
+    if mode == "sp":
+        return get_sp_group()
+    elif mode == "ulysses":
+        return get_sp_group().ulysses_group
+    elif mode == "ring":
+        return get_sp_group().ring_group
+    elif mode == "world":
+        # the whole single-replica DiT (all GPUs), regardless of tp/sp/cfg.
+        return get_world_group()
+    return get_tp_group()
+
+
+# Folding pays off only for wide encoders: measured ~-22% encode latency for
+# T5-XXL (hidden 4096) and larger for Mistral-24B (hidden 5120), but a net loss
+# for narrower ones (Qwen3 hidden 2560, CLIP 512) whose per-layer all_reduce
+# dominates the sharded compute. Decided on the real (post-load) hidden size.
+FOLD_MIN_HIDDEN_SIZE = 4096
+
+
+def _encoder_dims(config: EncoderConfig):
+    """Best-effort (hidden, attention_heads, mlp_intermediate) from a config,
+    spelled differently across families (hidden_size/d_model, num_heads, d_ff)."""
+
+    def first(names):
+        for name in names:
+            value = getattr(config, name, None)
+            if isinstance(value, int) and value > 0:
+                return value
+        return None
+
+    return (
+        first(("hidden_size", "d_model")),
+        first(("num_attention_heads", "num_heads", "n_heads")),
+        first(("intermediate_size", "d_ff", "ffn_dim")),
+    )
+
+
+def encoder_folding_worthwhile(config: EncoderConfig, group_size: int) -> bool:
+    """Fold only encoders wide enough to benefit whose heads and MLP divide the
+    fold group. Size-based (not per-architecture), so the same encoder family at
+    different parameter counts is handled correctly."""
+    hidden, heads, inter = _encoder_dims(config)
+    return (
+        group_size > 1
+        and hidden is not None
+        and hidden >= FOLD_MIN_HIDDEN_SIZE
+        and heads is not None
+        and heads % group_size == 0
+        and inter is not None
+        and inter % group_size == 0
+    )
+
+
+def finalize_encoder_folding(config: EncoderConfig) -> None:
+    """Loader hook: call after the encoder's real dims are populated
+    (update_model_arch) and before construction. adjust_pipeline_config proposes
+    a fold group from the parallelism alone; here we keep it only if the encoder
+    is actually worth folding at its real size, otherwise fall back to
+    replicated by clearing the mode.
+    """
+    if config.parallel_folding_mode is None:
+        return
+    group_size = getattr(get_folding_tp_group(config), "world_size", 1)
+    if not encoder_folding_worthwhile(config, group_size):
+        config.parallel_folding_mode = None
+
+
+class TextEncoder(nn.Module, ABC, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "layers",
+        "encoder.block",
+        "text_model.encoder.layers",
+        "model.language_model.layers",
+    ]
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
+    _stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
+    _supported_attention_backends: set[AttentionBackendEnum] = (
+        TextEncoderConfig()._supported_attention_backends
+    )
+
+    def __init__(self, config: TextEncoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        self._fsdp_shard_conditions = config.arch_config._fsdp_shard_conditions
+        self._stacked_params_mapping = config.arch_config.stacked_params_mapping
+        if not self.supported_attention_backends:
+            raise ValueError(
+                f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
+            )
+
+    @abstractmethod
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        pass
+
+    @property
+    def supported_attention_backends(self) -> set[AttentionBackendEnum]:
+        return self._supported_attention_backends
+
+
+class ImageEncoder(nn.Module, ABC, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "layers",
+        "vision_model.encoder.layers",
+        "model.visual.blocks",
+    ]
+    _supported_attention_backends: set[AttentionBackendEnum] = (
+        ImageEncoderConfig()._supported_attention_backends
+    )
+
+    def __init__(self, config: ImageEncoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        if not self.supported_attention_backends:
+            raise ValueError(
+                f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
+            )
+
+    @abstractmethod
+    def forward(self, pixel_values: torch.Tensor, **kwargs) -> BaseEncoderOutput:
+        pass
+
+    @property
+    def supported_attention_backends(self) -> set[AttentionBackendEnum]:
+        return self._supported_attention_backends

python/sglang/multimodal_gen/runtime/models/encoders/bert.py

@@ -0,0 +1,46 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# type: ignore
+import os
+
+import torch
+import torch.nn as nn
+from transformers import BertModel, BertTokenizer
+
+
+class HunyuanClip(nn.Module):
+    """
+    Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
+    hunyuan's clip used BertModel and BertTokenizer, so we copy it.
+    """
+
+    def __init__(self, model_dir, max_length=77):
+        super().__init__()
+
+        self.max_length = max_length
+        self.tokenizer = BertTokenizer.from_pretrained(
+            os.path.join(model_dir, "tokenizer")
+        )
+        self.text_encoder = BertModel.from_pretrained(
+            os.path.join(model_dir, "clip_text_encoder")
+        )
+
+    @torch.no_grad
+    def forward(self, prompts, with_mask=True):
+        self.device = next(self.text_encoder.parameters()).device
+        text_inputs = self.tokenizer(
+            prompts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        prompt_embeds = self.text_encoder(
+            text_inputs.input_ids.to(self.device),
+            attention_mask=(
+                text_inputs.attention_mask.to(self.device) if with_mask else None
+            ),
+        )
+        return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output

python/sglang/multimodal_gen/runtime/models/encoders/clip.py

@@ -0,0 +1,804 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/clip.py
+# Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py
+"""Minimal implementation of CLIPVisionModel intended to be only used
+within a vision language model."""
+
+from collections.abc import Iterable
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.configs.models.encoders import (
+    BaseEncoderOutput,
+    CLIPTextConfig,
+    CLIPVisionConfig,
+)
+from sglang.multimodal_gen.runtime.distributed import divide, get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+
+# TODO: support quantization
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+from sglang.multimodal_gen.runtime.models.encoders.base import ImageEncoder, TextEncoder
+from sglang.multimodal_gen.runtime.models.encoders.vision import (
+    resolve_visual_encoder_outputs,
+)
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py#L164 # noqa
+class CLIPVisionEmbeddings(nn.Module):
+
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+        assert self.image_size % self.patch_size == 0
+
+        self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches + 1
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+        self.register_buffer(
+            "position_ids",
+            torch.arange(self.num_positions).expand((1, -1)),
+            persistent=False,
+        )
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        batch_size = pixel_values.shape[0]
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(
+            pixel_values.to(dtype=target_dtype)
+        )  # shape = [*, width, grid, grid]
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+        class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+        embeddings = embeddings + self.position_embedding(self.position_ids)
+
+        return embeddings
+
+
+class CLIPTextEmbeddings(nn.Module):
+
+    def __init__(self, config: CLIPTextConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
+        self.position_embedding = nn.Embedding(
+            config.max_position_embeddings, embed_dim
+        )
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer(
+            "position_ids",
+            torch.arange(config.max_position_embeddings).expand((1, -1)),
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+    ) -> torch.Tensor:
+        if input_ids is not None:
+            seq_length = input_ids.shape[-1]
+        elif inputs_embeds is not None:
+            seq_length = inputs_embeds.shape[-2]
+        else:
+            raise ValueError("Either input_ids or inputs_embeds must be provided.")
+
+        max_position_embedding = self.position_embedding.weight.shape[0]
+
+        if seq_length > max_position_embedding:
+            raise ValueError(
+                f"Sequence length must be less than max_position_embeddings (got `sequence length`: "
+                f"{seq_length} and max_position_embeddings: {max_position_embedding}"
+            )
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.token_embedding(input_ids)
+
+        position_embeddings = self.position_embedding(position_ids)
+        embeddings = inputs_embeds + position_embeddings
+
+        return embeddings
+
+
+class CLIPAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig | CLIPTextConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                "embed_dim must be divisible by num_heads "
+                f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.embed_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.out_proj = RowParallelLinear(
+            input_size=self.embed_dim,
+            output_size=self.embed_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+        )
+
+        self.tp_size = get_tp_world_size()
+        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
+
+        self.attn = LocalAttention(
+            self.num_heads_per_partition,
+            self.head_dim,
+            self.num_heads_per_partition,
+            softmax_scale=self.scale,
+            causal=True,
+            supported_attention_backends=config._supported_attention_backends,
+        )
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return (
+            tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        """Input shape: Batch x Time x Channel"""
+
+        qkv_states, _ = self.qkv_proj(hidden_states)
+        query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
+        # use flash_attn_func
+        query_states = query_states.reshape(
+            query_states.shape[0],
+            query_states.shape[1],
+            self.num_heads_per_partition,
+            self.head_dim,
+        )
+        key_states = key_states.reshape(
+            key_states.shape[0],
+            key_states.shape[1],
+            self.num_heads_per_partition,
+            self.head_dim,
+        )
+        value_states = value_states.reshape(
+            value_states.shape[0],
+            value_states.shape[1],
+            self.num_heads_per_partition,
+            self.head_dim,
+        )
+
+        if self.attn.backend == AttentionBackendEnum.TORCH_SDPA:
+            query_states = query_states.transpose(1, 2)  # [B, H, S, D]
+            key_states = key_states.transpose(1, 2)
+            value_states = value_states.transpose(1, 2)
+
+            if (
+                current_platform.is_rocm()
+                or current_platform.is_musa()
+                or current_platform.is_xpu()
+            ):
+                # ROCm: Using both is_causal=True and attn_mask causes NaN.
+                # Use is_causal=True alone (padding mask not needed for CLIP
+                # since pooler_output comes from EOS token before padding).
+                # XXX (MUSA): Torch SDPA on MUSA currently does not support
+                # using both `attn_mask` and `is_causal=True` simultaneously.
+                attn_output = torch.nn.functional.scaled_dot_product_attention(
+                    query_states,
+                    key_states,
+                    value_states,
+                    attn_mask=None,
+                    is_causal=True,
+                    scale=self.scale,
+                )
+            else:
+                if attention_mask is not None:
+                    # SDPA requires [B, 1, 1, S] or [B, S, S] format mask
+                    if attention_mask.dim() == 2:
+                        attn_mask = attention_mask[:, None, None, :].to(
+                            dtype=query_states.dtype
+                        )
+                        attn_mask = (1.0 - attn_mask) * torch.finfo(
+                            query_states.dtype
+                        ).min
+                    else:
+                        attn_mask = attention_mask
+                else:
+                    attn_mask = None
+
+                attn_output = torch.nn.functional.scaled_dot_product_attention(
+                    query_states,
+                    key_states,
+                    value_states,
+                    attn_mask=attn_mask,
+                    is_causal=attention_mask is None,
+                    scale=self.scale,
+                )
+            attn_output = attn_output.transpose(1, 2)
+        else:
+            # Use LocalAttention (doesn't support attention_mask, but maintains compatibility)
+            attn_output = self.attn(query_states, key_states, value_states)
+
+        attn_output = attn_output.reshape(
+            attn_output.shape[0],
+            attn_output.shape[1],
+            self.num_heads_per_partition * self.head_dim,
+        )
+        attn_output, _ = self.out_proj(attn_output)
+
+        return attn_output, None
+
+
+class CLIPMLP(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig | CLIPTextConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.activation_fn = get_act_fn(config.hidden_act)
+        self.fc1 = ColumnParallelLinear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc1",
+        )
+        self.fc2 = RowParallelLinear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc2",
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states, _ = self.fc2(hidden_states)
+
+        return hidden_states
+
+
+class CLIPEncoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPTextConfig | CLIPVisionConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = CLIPAttention(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = CLIPMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp")
+        self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class CLIPEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self
+    attention layers. Each layer is a [`CLIPEncoderLayer`].
+
+    Args:
+        config: CLIPConfig
+    """
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig | CLIPTextConfig,
+        quant_config: QuantizationConfig | None = None,
+        num_hidden_layers_override: int | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+
+        if num_hidden_layers_override is None:
+            num_hidden_layers = config.num_hidden_layers
+        else:
+            num_hidden_layers = num_hidden_layers_override
+        self.layers = nn.ModuleList(
+            [
+                CLIPEncoderLayer(
+                    config=config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{layer_idx}",
+                )
+                for layer_idx in range(num_hidden_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        return_all_hidden_states: bool,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        hidden_states_pool = [inputs_embeds]
+        hidden_states = inputs_embeds
+
+        for idx, encoder_layer in enumerate(self.layers):
+            hidden_states = encoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+            )
+            if return_all_hidden_states:
+                hidden_states_pool.append(hidden_states)
+        # If we have multiple feature sample layers, we return all hidden
+        # states in order and grab the ones we need by index.
+        if return_all_hidden_states:
+            return hidden_states_pool
+        return [hidden_states]
+
+
+class CLIPTextTransformer(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPTextConfig,
+        quant_config: QuantizationConfig | None = None,
+        num_hidden_layers_override: int | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = CLIPTextEmbeddings(config)
+
+        self.encoder = CLIPEncoder(
+            config,
+            quant_config=quant_config,
+            num_hidden_layers_override=num_hidden_layers_override,
+            prefix=prefix,
+        )
+
+        self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+        # For `pooled_output` computation
+        self.eos_token_id = config.eos_token_id
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+
+        if input_ids is None:
+            raise ValueError("You have to specify input_ids")
+
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
+
+        # CLIP's text model uses causal mask, prepare it here.
+        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
+        # causal_attention_mask = _create_4d_causal_attention_mask(
+        #     input_shape, hidden_states.dtype, device=hidden_states.device
+        # )
+
+        # # expand attention_mask
+        # if attention_mask is not None and not self._use_flash_attention_2:
+        #     raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer")
+        #     # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        #     attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            return_all_hidden_states=output_hidden_states,
+            attention_mask=attention_mask,
+        )
+
+        last_hidden_state = encoder_outputs[-1]
+        last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+        if self.eos_token_id == 2:
+            # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
+            # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added
+            # ------------------------------------------------------------
+            # text_embeds.shape = [batch_size, sequence_length, transformer.width]
+            # take features from the eot embedding (eot_token is the highest number in each sequence)
+            # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+            pooled_output = last_hidden_state[
+                torch.arange(
+                    last_hidden_state.shape[0], device=last_hidden_state.device
+                ),
+                input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(
+                    dim=-1
+                ),
+            ]
+        else:
+            # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
+            pooled_output = last_hidden_state[
+                torch.arange(
+                    last_hidden_state.shape[0], device=last_hidden_state.device
+                ),
+                # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
+                # Note: we assume each sequence (along batch dim.) contains an  `eos_token_id` (e.g. prepared by the tokenizer)
+                (
+                    input_ids.to(dtype=torch.int, device=last_hidden_state.device)
+                    == self.eos_token_id
+                )
+                .int()
+                .argmax(dim=-1),
+            ]
+
+        return BaseEncoderOutput(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs,
+            # attentions=encoder_outputs.attentions,
+        )
+
+
+class CLIPTextModel(TextEncoder):
+
+    def __init__(
+        self,
+        config: CLIPTextConfig,
+    ) -> None:
+        super().__init__(config)
+        self.text_model = CLIPTextTransformer(
+            config=config, quant_config=config.quant_config, prefix=config.prefix
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+
+        outputs: BaseEncoderOutput = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_hidden_states=output_hidden_states,
+        )
+        return outputs
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+
+        # Define mapping for stacked parameters
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            # Handle q_proj, k_proj, v_proj -> qkv_proj mapping
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name in name:
+                    # Replace the weight name with the parameter name
+                    model_param_name = name.replace(weight_name, param_name)
+
+                    if model_param_name in params_dict:
+                        param = params_dict[model_param_name]
+                        weight_loader = param.weight_loader
+                        weight_loader(param, loaded_weight, shard_id)
+                        loaded_params.add(model_param_name)
+                    break
+            else:
+                # Use default weight loader for all other parameters
+                if name in params_dict:
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(name)
+
+        return loaded_params
+
+
+class CLIPTextModelWithProjection(CLIPTextModel):
+    """
+    CLIP text encoder with projection head for pooled_output.
+    """
+
+    def __init__(
+        self,
+        config: CLIPTextConfig,
+    ) -> None:
+        super().__init__(config)
+        self.text_projection = nn.Linear(
+            config.hidden_size, config.projection_dim, bias=False
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        outputs: BaseEncoderOutput = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_hidden_states=output_hidden_states,
+        )
+
+        pooled_output = outputs.pooler_output
+        if pooled_output is not None:
+            pooled_output = self.text_projection(pooled_output)
+
+        return BaseEncoderOutput(
+            last_hidden_state=outputs.last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class CLIPVisionTransformer(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig,
+        quant_config: QuantizationConfig | None = None,
+        num_hidden_layers_override: int | None = None,
+        require_post_norm: bool | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = CLIPVisionEmbeddings(config)
+
+        # NOTE: This typo of "layrnorm" is not fixed on purpose to match
+        # the original transformers code and name of the model weights.
+        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+        self.encoder = CLIPEncoder(
+            config=config,
+            quant_config=quant_config,
+            num_hidden_layers_override=num_hidden_layers_override,
+            prefix=f"{prefix}.encoder",
+        )
+
+        num_hidden_layers = config.num_hidden_layers
+        if len(self.encoder.layers) > config.num_hidden_layers:
+            raise ValueError(
+                f"The original encoder only has {num_hidden_layers} "
+                f"layers, but you requested {len(self.encoder.layers)} layers."
+            )
+
+        # If possible, skip post_layernorm to conserve memory
+        if require_post_norm is None:
+            require_post_norm = len(self.encoder.layers) == num_hidden_layers
+
+        if require_post_norm:
+            self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        else:
+            self.post_layernorm = None
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        feature_sample_layers: list[int] | None = None,
+    ) -> BaseEncoderOutput:
+
+        hidden_states = self.embeddings(pixel_values)
+        hidden_states = self.pre_layrnorm(hidden_states)
+
+        return_all_hidden_states = output_hidden_states or (
+            feature_sample_layers is not None
+        )
+
+        # Produces either the last layer output or all of the hidden states,
+        # depending on if we have feature_sample_layers or not
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            return_all_hidden_states=return_all_hidden_states,
+        )
+
+        if not return_all_hidden_states:
+            encoder_outputs = encoder_outputs[0]
+
+            # Handle post-norm (if applicable) and stacks feature layers if needed
+            encoder_outputs = resolve_visual_encoder_outputs(
+                encoder_outputs,
+                feature_sample_layers,
+                self.post_layernorm,
+                self.config.num_hidden_layers,
+            )
+
+        if return_all_hidden_states:
+            return BaseEncoderOutput(hidden_states=encoder_outputs)
+
+        return BaseEncoderOutput(last_hidden_state=encoder_outputs)
+
+
+class CLIPVisionModel(ImageEncoder):
+    config_class = CLIPVisionConfig
+    main_input_name = "pixel_values"
+    packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
+
+    def __init__(self, config: CLIPVisionConfig) -> None:
+        super().__init__(config)
+        self.vision_model = CLIPVisionTransformer(
+            config=config,
+            quant_config=config.quant_config,
+            num_hidden_layers_override=config.num_hidden_layers_override,
+            require_post_norm=config.require_post_norm,
+            prefix=f"{config.prefix}.vision_model",
+        )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        feature_sample_layers: list[int] | None = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        base_encoder_output = self.vision_model(
+            pixel_values,
+            output_hidden_states=output_hidden_states,
+            feature_sample_layers=feature_sample_layers,
+        )
+
+        return base_encoder_output
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    # (TODO) Add prefix argument for filtering out weights to be loaded
+    #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        layer_count = len(self.vision_model.encoder.layers)
+
+        for name, loaded_weight in weights:
+            if name.startswith("visual_projection"):
+                continue
+            # post_layernorm is not needed in CLIPVisionModel
+            if (
+                name.startswith("vision_model.post_layernorm")
+                and self.vision_model.post_layernorm is None
+            ):
+                continue
+
+            # omit layers when num_hidden_layers_override is set
+            if name.startswith("vision_model.encoder.layers"):
+                layer_idx = int(name.split(".")[3])
+                if layer_idx >= layer_count:
+                    continue
+
+            for (
+                param_name,
+                weight_name,
+                shard_id,
+            ) in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class BertModel(CLIPTextModel):
+    pass
+
+
+EntryClass = [CLIPTextModel, CLIPTextModelWithProjection, CLIPVisionModel]

python/sglang/multimodal_gen/runtime/models/encoders/gemma2.py

@@ -0,0 +1,447 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Gemma2 2B text encoder for SANA.
+#
+# This is a decoder-only language model used as a text encoder: we feed
+# in tokenized text and extract the final hidden states (not logits) as
+# the conditioning signal for SANA's cross-attention layers.
+#
+# Architecture follows google/gemma-2-2b-it:
+#   - 26 layers, alternating global / sliding-window attention
+#   - GQA with 8 query heads, 4 KV heads, head_dim=256
+#   - Pre/post attention + pre/post feedforward LayerNorm (Gemma2-style)
+#   - GeGLU activation (gelu_pytorch_tanh)
+#
+# Adapted from the Gemma3 text model implementation in this codebase.
+
+import logging
+from typing import Any, Iterable
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Gemma2RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float())
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+
+class Gemma2MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "gelu_pytorch_tanh":
+            raise ValueError(
+                "Gemma2 uses `gelu_pytorch_tanh` as the hidden activation. "
+                f"Got: {hidden_act}"
+            )
+        self.act_fn = GeluAndMul(approximate="tanh")
+
+    def forward(self, x):
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Gemma2Attention(nn.Module):
+    def __init__(
+        self,
+        layer_id: int,
+        config: Gemma2Config,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.layer_id = layer_id
+        self.hidden_size = hidden_size
+        tp_size = get_tp_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        arch = config.arch_config
+        self.head_dim = arch.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = arch.query_pre_attn_scalar**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=arch.attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=arch.attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # Gemma2 interleaves global (even layers) and sliding-window (odd layers)
+        # attention. This pattern reduces memory for long sequences while
+        # maintaining global context every other layer.
+        self.is_sliding = (layer_id % 2) == 1
+        if self.is_sliding:
+            self.sliding_window = arch.sliding_window
+        else:
+            self.sliding_window = None
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=arch.max_position_embeddings,
+            base=arch.rope_theta,
+            is_neox_style=True,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        batch_size, seq_len, _ = q.shape
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        query = q.transpose(1, 2)
+        key = k.transpose(1, 2)
+        value = v.transpose(1, 2)
+
+        attn_mask = torch.tril(
+            torch.ones(
+                (seq_len, seq_len), device=hidden_states.device, dtype=torch.bool
+            )
+        )
+        if self.is_sliding and self.sliding_window is not None:
+            idx = torch.arange(seq_len, device=hidden_states.device)
+            dist = idx[None, :] - idx[:, None]
+            too_far = dist > self.sliding_window
+            attn_mask = attn_mask.masked_fill(too_far, False)
+
+        if attention_mask is not None:
+            attn_mask = attn_mask[None, None, :, :].expand(
+                batch_size, 1, seq_len, seq_len
+            )
+            attn_mask = attn_mask & attention_mask.to(torch.bool)[:, None, None, :]
+
+        attn_kwargs = {
+            "attn_mask": attn_mask,
+            "dropout_p": 0.0,
+            "is_causal": False,
+            "scale": self.scaling,
+        }
+        if query.shape[1] != key.shape[1]:
+            attn_kwargs["enable_gqa"] = True
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, **attn_kwargs
+        )
+
+        # NOTE: Gemma2 specifies attn_logit_softcapping (tanh(logits/cap)*cap) but
+        # PyTorch's scaled_dot_product_attention does not support it natively.
+        # For short text-encoder sequences (~300 tokens), the quality impact is
+        # negligible. A custom attention kernel would be needed for full fidelity.
+
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(
+            batch_size, seq_len, self.num_heads * self.head_dim
+        )
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Gemma2DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        layer_id: int,
+        config: Gemma2Config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        arch = config.arch_config
+        self.hidden_size = arch.hidden_size
+        self.self_attn = Gemma2Attention(
+            layer_id=layer_id,
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=arch.num_attention_heads,
+            num_kv_heads=arch.num_key_value_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = Gemma2MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=arch.intermediate_size,
+            hidden_act=arch.hidden_activation,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = Gemma2RMSNorm(self.hidden_size, eps=arch.rms_norm_eps)
+        self.post_attention_layernorm = Gemma2RMSNorm(
+            self.hidden_size, eps=arch.rms_norm_eps
+        )
+        self.pre_feedforward_layernorm = Gemma2RMSNorm(
+            self.hidden_size, eps=arch.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = Gemma2RMSNorm(
+            self.hidden_size, eps=arch.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(positions, hidden_states, attention_mask)
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.pre_feedforward_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.post_feedforward_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class Gemma2Model(nn.Module, LayerwiseOffloadableModuleMixin):
+    """Gemma2 text encoder model for SANA pipeline."""
+
+    _fsdp_shard_conditions = []
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["layers"]
+
+    def __init__(self, config: Gemma2Config, **kwargs):
+        super().__init__()
+        self.config = config
+        arch = config.arch_config
+        self.quant_config = None
+
+        self.vocab_size = arch.vocab_size
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            arch.hidden_size,
+            org_num_embeddings=arch.vocab_size,
+            quant_config=self.quant_config,
+        )
+        self.embed_scale = arch.hidden_size**0.5
+
+        self.layers = nn.ModuleList(
+            [
+                Gemma2DecoderLayer(
+                    layer_id=i,
+                    config=config,
+                    quant_config=self.quant_config,
+                    prefix=f"model.layers.{i}",
+                )
+                for i in range(arch.num_hidden_layers)
+            ]
+        )
+
+        self.norm = Gemma2RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids) * self.embed_scale
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else getattr(self.config.arch_config, "output_hidden_states", False)
+        )
+
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, hidden_states.shape[1], device=hidden_states.device
+            ).unsqueeze(0)
+
+        all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
+
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                all_hidden_states += (hidden_states,)
+
+            hidden_states = layer(position_ids, hidden_states, attention_mask)
+
+        hidden_states = self.norm(hidden_states)
+
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states,)
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        stacked_params_mapping = getattr(
+            self.config.arch_config, "stacked_params_mapping", None
+        )
+        if stacked_params_mapping is None:
+            stacked_params_mapping = [
+                (".qkv_proj", ".q_proj", "q"),
+                (".qkv_proj", ".k_proj", "k"),
+                (".qkv_proj", ".v_proj", "v"),
+                (".gate_up_proj", ".gate_proj", "0"),
+                (".gate_up_proj", ".up_proj", "1"),
+            ]
+
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            # HF Gemma2Model stores weights as model.layers.X... / model.embed_tokens...
+            # Strip "model." prefix if present to match our naming
+            if name.startswith("model."):
+                name = name[len("model.") :]
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                self._load_with_shard_id(weight_loader, param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(name)
+        return loaded_params
+
+    @staticmethod
+    def _load_with_shard_id(weight_loader, param, loaded_weight, shard_id):
+        try:
+            weight_loader(param, loaded_weight, shard_id)
+            return
+        except (AssertionError, TypeError):
+            pass
+
+        if isinstance(shard_id, str):
+            mapping = {"q": 0, "k": 1, "v": 2}
+            if shard_id in mapping:
+                weight_loader(param, loaded_weight, mapping[shard_id])
+                return
+            if shard_id.isdigit():
+                weight_loader(param, loaded_weight, int(shard_id))
+                return
+        elif isinstance(shard_id, int):
+            mapping = {0: "q", 1: "k", 2: "v"}
+            if shard_id in mapping:
+                weight_loader(param, loaded_weight, mapping[shard_id])
+                return
+
+        raise TypeError(
+            f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader}"
+        )
+
+
+EntryClass = Gemma2Model

python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py

@@ -0,0 +1,287 @@
+# Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from transformers import (
+    CLIPVisionConfig,
+    CLIPVisionModelWithProjection,
+    Dinov2Config,
+    Dinov2Model,
+)
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    return np.concatenate([emb_sin, emb_cos], axis=1)
+
+
+class ImageEncoder(nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "model.encoder.layer",
+        "model.vision_model.encoder.layers",
+    ]
+    MODEL_CLASS = None
+    MODEL_CONFIG_CLASS = None
+    mean = []
+    std = []
+
+    def __init__(
+        self,
+        version=None,
+        config=None,
+        use_cls_token=True,
+        image_size=224,
+        **kwargs,
+    ):
+        super().__init__()
+
+        if config is None:
+            self.model = self.MODEL_CLASS.from_pretrained(version)
+        else:
+            self.model = self.MODEL_CLASS(self.MODEL_CONFIG_CLASS.from_dict(config))
+        self.model.eval()
+        self.model.requires_grad_(False)
+        self.use_cls_token = use_cls_token
+        self.size = image_size // 14
+        self.num_patches = (image_size // 14) ** 2
+        if self.use_cls_token:
+            self.num_patches += 1
+
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize(
+                    image_size, transforms.InterpolationMode.BILINEAR, antialias=True
+                ),
+                transforms.CenterCrop(image_size),
+                transforms.Normalize(
+                    mean=self.mean,
+                    std=self.std,
+                ),
+            ]
+        )
+
+    def forward(self, image, mask=None, value_range=(-1, 1), **kwargs):
+        if value_range is not None:
+            low, high = value_range
+            image = (image - low) / (high - low)
+
+        image = image.to(self.model.device, dtype=self.model.dtype)
+        inputs = self.transform(image)
+        outputs = self.model(inputs)
+
+        last_hidden_state = outputs.last_hidden_state
+        if not self.use_cls_token:
+            last_hidden_state = last_hidden_state[:, 1:, :]
+
+        return last_hidden_state
+
+    def unconditional_embedding(self, batch_size, **kwargs):
+        device = next(self.model.parameters()).device
+        dtype = next(self.model.parameters()).dtype
+        zero = torch.zeros(
+            batch_size,
+            self.num_patches,
+            self.model.config.hidden_size,
+            device=device,
+            dtype=dtype,
+        )
+
+        return zero
+
+
+class CLIPImageEncoder(ImageEncoder):
+    MODEL_CLASS = CLIPVisionModelWithProjection
+    MODEL_CONFIG_CLASS = CLIPVisionConfig
+    mean = [0.48145466, 0.4578275, 0.40821073]
+    std = [0.26862954, 0.26130258, 0.27577711]
+
+
+class DinoImageEncoder(ImageEncoder):
+    MODEL_CLASS = Dinov2Model
+    MODEL_CONFIG_CLASS = Dinov2Config
+    mean = [0.485, 0.456, 0.406]
+    std = [0.229, 0.224, 0.225]
+
+
+class DinoImageEncoderMV(DinoImageEncoder):
+    _aliases = [
+        "hy3dshape.models.conditioner.DinoImageEncoderMV",
+    ]
+
+    def __init__(
+        self,
+        version=None,
+        config=None,
+        use_cls_token=True,
+        image_size=224,
+        view_num=4,
+        **kwargs,
+    ):
+        super().__init__(version, config, use_cls_token, image_size, **kwargs)
+        self.view_num = view_num
+        self.num_patches = self.num_patches
+        pos = np.arange(self.view_num, dtype=np.float32)
+        view_embedding = torch.from_numpy(
+            get_1d_sincos_pos_embed_from_grid(self.model.config.hidden_size, pos)
+        ).float()
+
+        view_embedding = view_embedding.unsqueeze(1).repeat(1, self.num_patches, 1)
+        self.view_embed = view_embedding.unsqueeze(0)
+
+    def forward(self, image, mask=None, value_range=(-1, 1), view_idxs=None, **kwargs):
+        if value_range is not None:
+            low, high = value_range
+            image = (image - low) / (high - low)
+
+        image = image.to(self.model.device, dtype=self.model.dtype)
+
+        bs, num_views, c, h, w = image.shape
+        image = image.view(bs * num_views, c, h, w)
+
+        inputs = self.transform(image)
+        outputs = self.model(inputs)
+
+        last_hidden_state = outputs.last_hidden_state
+        last_hidden_state = last_hidden_state.view(
+            bs, num_views, last_hidden_state.shape[-2], last_hidden_state.shape[-1]
+        )
+
+        view_embedding = self.view_embed.to(last_hidden_state.dtype).to(
+            last_hidden_state.device
+        )
+        if view_idxs is not None:
+            assert len(view_idxs) == bs
+            view_embeddings = []
+            for i in range(bs):
+                view_idx = view_idxs[i]
+                assert num_views == len(view_idx)
+                view_embeddings.append(self.view_embed[:, view_idx, ...])
+            view_embedding = (
+                torch.cat(view_embeddings, 0)
+                .to(last_hidden_state.dtype)
+                .to(last_hidden_state.device)
+            )
+
+        if num_views != self.view_num:
+            view_embedding = view_embedding[:, :num_views, ...]
+        last_hidden_state = last_hidden_state + view_embedding
+        last_hidden_state = last_hidden_state.view(
+            bs, num_views * last_hidden_state.shape[-2], last_hidden_state.shape[-1]
+        )
+        return last_hidden_state
+
+    def unconditional_embedding(self, batch_size, view_idxs, **kwargs):
+        device = next(self.model.parameters()).device
+        dtype = next(self.model.parameters()).dtype
+        zero = torch.zeros(
+            batch_size,
+            self.num_patches * len(view_idxs[0]),
+            self.model.config.hidden_size,
+            device=device,
+            dtype=dtype,
+        )
+        return zero
+
+
+def build_image_encoder(config):
+    if config["type"] == "CLIPImageEncoder":
+        return CLIPImageEncoder(**config["kwargs"])
+    elif config["type"] == "DinoImageEncoder":
+        return DinoImageEncoder(**config["kwargs"])
+    elif config["type"] == "DinoImageEncoderMV":
+        return DinoImageEncoderMV(**config["kwargs"])
+    else:
+        raise ValueError(f'Unknown image encoder type: {config["type"]}')
+
+
+class DualImageEncoder(nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "main_image_encoder.model.encoder.layer",
+        "main_image_encoder.model.vision_model.encoder.layers",
+        "additional_image_encoder.model.encoder.layer",
+        "additional_image_encoder.model.vision_model.encoder.layers",
+    ]
+
+    def __init__(
+        self,
+        main_image_encoder,
+        additional_image_encoder,
+    ):
+        super().__init__()
+        self.main_image_encoder = build_image_encoder(main_image_encoder)
+        self.additional_image_encoder = build_image_encoder(additional_image_encoder)
+
+    def forward(self, image, mask=None, **kwargs):
+        outputs = {
+            "main": self.main_image_encoder(image, mask=mask, **kwargs),
+            "additional": self.additional_image_encoder(image, mask=mask, **kwargs),
+        }
+        return outputs
+
+    def unconditional_embedding(self, batch_size, **kwargs):
+        outputs = {
+            "main": self.main_image_encoder.unconditional_embedding(
+                batch_size, **kwargs
+            ),
+            "additional": self.additional_image_encoder.unconditional_embedding(
+                batch_size, **kwargs
+            ),
+        }
+        return outputs
+
+
+class SingleImageEncoder(nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "main_image_encoder.model.encoder.layer",
+        "main_image_encoder.model.vision_model.encoder.layers",
+    ]
+
+    def __init__(
+        self,
+        main_image_encoder,
+    ):
+        super().__init__()
+        self.main_image_encoder = build_image_encoder(main_image_encoder)
+
+    def forward(self, image, mask=None, **kwargs):
+        outputs = {
+            "main": self.main_image_encoder(image, mask=mask, **kwargs),
+        }
+        return outputs
+
+    def unconditional_embedding(self, batch_size, **kwargs):
+        outputs = {
+            "main": self.main_image_encoder.unconditional_embedding(
+                batch_size, **kwargs
+            ),
+        }
+        return outputs
+
+
+# Entry class for model registry
+EntryClass = [
+    SingleImageEncoder,
+    DualImageEncoder,
+    DinoImageEncoder,
+    DinoImageEncoderMV,
+]

python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Iterable
+from typing import Tuple
+
+import torch
+from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig
+
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.ideogram import (
+    Ideogram4TextEncoderConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
+    BitsAndBytesConfig,
+    attach_bitsandbytes_4bit_quant_states,
+    build_bitsandbytes_4bit_quant_states,
+    is_bitsandbytes_4bit_state_name,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
+from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextModel
+
+
+class IdeogramQwen3VLTextEncoder(TextEncoder):
+    """Language-only Qwen3-VL text encoder stored inside Ideogram checkpoints."""
+
+    _activation_layers = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
+
+    def __init__(self, config: Ideogram4TextEncoderConfig) -> None:
+        super().__init__(config)
+        arch_config = config.arch_config
+        text_config = getattr(arch_config, "text_config")
+        if isinstance(text_config, dict):
+            text_config = Qwen3VLTextConfig(**text_config)
+        self._uses_bitsandbytes_4bit = getattr(
+            arch_config, "ideogram_bnb_4bit_weight_only", False
+        )
+        self._uses_weight_only_fp8 = getattr(
+            arch_config, "ideogram_fp8_weight_only", False
+        )
+        quant_config = None
+        if self._uses_bitsandbytes_4bit:
+            source_quant_config = getattr(arch_config, "quantization_config")
+            if isinstance(source_quant_config, dict):
+                quant_config_dict = source_quant_config
+            else:
+                quant_config_dict = source_quant_config.to_dict()
+            quant_config = BitsAndBytesConfig.from_config(quant_config_dict)
+        self.language_model = Qwen3VLTextModel(
+            text_config,
+            quant_config=quant_config,
+            use_weight_only_fp8=self._uses_weight_only_fp8,
+            use_tensor_parallel=True,
+        )
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        outputs = self.language_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            use_cache=False,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        return BaseEncoderOutput(
+            last_hidden_state=outputs.last_hidden_state,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def encode_ideogram_features(
+        self,
+        token_ids: torch.Tensor,
+        text_position_ids: torch.Tensor,
+        indicator: torch.Tensor,
+        llm_token_indicator: int,
+    ) -> torch.Tensor:
+        batch_size, seq_len = token_ids.shape
+        hidden_size = self.language_model.config.hidden_size
+        out_dim = hidden_size * len(self._activation_layers)
+        features = torch.zeros(
+            batch_size,
+            seq_len,
+            out_dim,
+            dtype=torch.float32,
+            device=token_ids.device,
+        )
+        for batch_idx in range(batch_size):
+            text_mask = indicator[batch_idx] == llm_token_indicator
+            cur_token_ids = token_ids[batch_idx, text_mask].unsqueeze(0)
+            if cur_token_ids.numel() == 0:
+                continue
+            pos_2d = text_position_ids[batch_idx, text_mask, 0].unsqueeze(0)
+            position_ids = pos_2d[None, ...].expand(4, 1, -1)
+            attention_mask = torch.ones_like(cur_token_ids)
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                outputs = self.forward(
+                    input_ids=cur_token_ids,
+                    position_ids=position_ids,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                )
+            assert outputs.hidden_states is not None
+            selected = [outputs.hidden_states[i] for i in self._activation_layers]
+            stacked = torch.stack(selected, dim=0).permute(1, 2, 3, 0)
+            features[batch_idx, text_mask] = stacked.reshape(
+                1, cur_token_ids.shape[1], -1
+            )[0].to(torch.float32)
+        return features
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        if self._uses_bitsandbytes_4bit:
+            return self._load_bitsandbytes_4bit_weights(weights)
+
+        loaded_params: set[str] = set()
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        for name, loaded_weight in weights:
+            if name.startswith("visual."):
+                continue
+            if "rotary_emb.inv_freq" in name:
+                continue
+            param = params_dict.get(name)
+            if param is None:
+                raise KeyError(
+                    f"Unexpected weight name while loading Ideogram text encoder: {name}"
+                )
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight.to(param.dtype))
+            loaded_params.add(name)
+        return loaded_params
+
+    def _load_bitsandbytes_4bit_weights(
+        self, weights: Iterable[Tuple[str, torch.Tensor]]
+    ):
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        raw_quant_state: dict[str, torch.Tensor] = {}
+        normal_weight_names: list[str] = []
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if is_bitsandbytes_4bit_state_name(name):
+                if "quant_state.bitsandbytes" in name:
+                    loaded_weight = loaded_weight.cpu().data
+                raw_quant_state[name] = loaded_weight
+                continue
+            if name.startswith("visual."):
+                continue
+            if "rotary_emb.inv_freq" in name:
+                continue
+            param = params_dict.get(name)
+            if param is None:
+                raise KeyError(
+                    f"Unexpected weight name while loading Ideogram text encoder: {name}"
+                )
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight.to(param.dtype))
+            normal_weight_names.append(name)
+            loaded_params.add(name)
+
+        quant_states = build_bitsandbytes_4bit_quant_states(
+            normal_weight_names,
+            raw_quant_state,
+            next(self.parameters()).device,
+        )
+        attach_bitsandbytes_4bit_quant_states(params_dict, quant_states)
+        quantized_params_missing_state = [
+            name
+            for name, param in params_dict.items()
+            if getattr(param, "use_bitsandbytes_4bit", False)
+            and name not in quant_states
+        ]
+        if quantized_params_missing_state:
+            raise ValueError(
+                "Missing bitsandbytes quant_state for Ideogram text encoder weights: "
+                f"{quantized_params_missing_state[:8]}"
+            )
+        return loaded_params
+
+
+EntryClass = IdeogramQwen3VLTextEncoder

python/sglang/multimodal_gen/runtime/models/encoders/llama.py

@@ -0,0 +1,460 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/llama.py
+
+# Adapted from
+# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only LLaMA model compatible with HuggingFace weights."""
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from torch import nn
+
+# from ..utils import (extract_layer_index)
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput, LlamaConfig
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
+
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
+
+
+class LlamaMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            # output_size=intermediate_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. "
+                "Only silu is supported for now."
+            )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class LlamaAttention(nn.Module):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 10000,
+        rope_scaling: dict[str, Any] | None = None,
+        max_position_embeddings: int = 8192,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        bias_o_proj: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        # layer_idx = extract_layer_index(prefix)
+        self.hidden_size = hidden_size
+        tp_size = get_tp_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
+        self.head_dim = getattr(
+            config, "head_dim", self.hidden_size // self.total_num_heads
+        )
+        # Phi models introduced a partial_rotary_factor parameter in the config
+        partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)
+        self.rotary_dim = int(partial_rotary_factor * self.head_dim)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias_o_proj,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        is_neox_style = True
+        is_gguf = (
+            quant_config
+            and hasattr(quant_config, "get_name")
+            and quant_config.get_name() == "gguf"
+        )
+        if is_gguf and config.model_type == "llama":
+            is_neox_style = False
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.rotary_dim,
+            max_position=max_position_embeddings,
+            base=int(rope_theta),
+            rope_scaling=rope_scaling,
+            is_neox_style=is_neox_style,
+        )
+
+        self.attn = LocalAttention(
+            self.num_heads,
+            self.head_dim,
+            self.num_kv_heads,
+            softmax_scale=self.scaling,
+            causal=True,
+            supported_attention_backends=config._supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        # attn_output = self.attn(q, k, v)
+        # use flash_attn_func
+        # TODO (Attn abstraction and backend)
+        # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim)
+        batch_size = q.shape[0]
+        seq_len = q.shape[1]
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        # import pdb; pdb.set_trace()
+        # attn_output = flash_attn_varlen_func(q, k, v, softmax_scale=self.scaling, causal=True)
+        attn_output = self.attn(q, k, v)
+        attn_output = attn_output.reshape(
+            batch_size, seq_len, self.num_heads * self.head_dim
+        )
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class LlamaDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = config.rope_parameters["rope_theta"]
+        rope_scaling = config.rope_parameters
+        if rope_scaling is not None and getattr(
+            config, "original_max_position_embeddings", None
+        ):
+            rope_scaling["original_max_position_embeddings"] = (
+                config.original_max_position_embeddings
+            )
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        # Support abacusai/Smaug-72B-v0.1 with attention_bias
+        # Support internlm/internlm-7b with bias
+        attention_bias = getattr(config, "attention_bias", False) or getattr(
+            config, "bias", False
+        )
+        bias_o_proj = attention_bias
+        # support internlm/internlm3-8b with qkv_bias
+        if hasattr(config, "qkv_bias"):
+            attention_bias = config.qkv_bias
+
+        self.self_attn = LlamaAttention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(
+                config, "num_key_value_heads", config.num_attention_heads
+            ),
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            bias_o_proj=bias_o_proj,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = LlamaMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            bias=getattr(config, "mlp_bias", False),
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class LlamaModel(TextEncoder):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+    ):
+        super().__init__(config)
+
+        self.config = config
+        self.quant_config = self.config.quant_config
+        if config.lora_config is not None:
+            max_loras = 1
+            lora_vocab_size = 1
+            if hasattr(config.lora_config, "max_loras"):
+                max_loras = config.lora_config.max_loras
+            if hasattr(config.lora_config, "lora_extra_vocab_size"):
+                lora_vocab_size = config.lora_config.lora_extra_vocab_size
+            lora_vocab = lora_vocab_size * max_loras
+        else:
+            lora_vocab = 0
+        self.vocab_size = config.vocab_size + lora_vocab
+        self.org_vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+            org_num_embeddings=config.vocab_size,
+            quant_config=config.quant_config,
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                LlamaDecoderLayer(
+                    config=config,
+                    quant_config=config.quant_config,
+                    prefix=f"{config.prefix}.layers.{i}",
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+        residual = None
+
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, hidden_states.shape[1], device=hidden_states.device
+            ).unsqueeze(0)
+
+        all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                # TODO
+                all_hidden_states += (
+                    (hidden_states,)
+                    if residual is None
+                    else (hidden_states + residual,)
+                )
+            hidden_states, residual = layer(position_ids, hidden_states, residual)
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        # add hidden states from the last decoder layer
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states,)
+
+        # TODO(will): maybe unify the output format with other models and use
+        # our own class
+        output = BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            # past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            # attentions=all_self_attns,
+        )
+
+        return output
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+            # if (self.quant_config is not None and
+            #     (scale_name := self.quant_config.get_cache_scale(name))):
+            #     # Loading kv cache quantization scales
+            #     param = params_dict[scale_name]
+            #     weight_loader = getattr(param, "weight_loader",
+            #                             default_weight_loader)
+            #     loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
+            #                      loaded_weight[0])
+            #     weight_loader(param, loaded_weight)
+            #     loaded_params.add(scale_name)
+            #     continue
+            if "scale" in name:
+                # Remapping the name of FP8 kv-scale.
+                kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
+                if kv_scale_name is None:
+                    continue
+                else:
+                    name = kv_scale_name
+            for (
+                param_name,
+                weight_name,
+                shard_id,
+            ) in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+EntryClass = LlamaModel

python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py

@@ -0,0 +1,630 @@
+# coding=utf-8
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import inspect
+from contextlib import nullcontext
+from typing import Iterable, Optional, Union
+
+import torch
+from torch import nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from transformers import Cache, DynamicCache, LlavaConfig, Mistral3Config, MistralConfig
+from transformers.activations import ACT2FN
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.mistral3.modeling_mistral3 import (
+    Mistral3CausalLMOutputWithPast,
+    Mistral3ModelOutputWithPast,
+)
+from transformers.models.mistral.modeling_mistral import (
+    MistralPreTrainedModel,
+    MistralRMSNorm,
+    MistralRotaryEmbedding,
+    apply_rotary_pos_emb,
+)
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_tp_world_size,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+_CREATE_CAUSAL_MASK_ARG = (
+    "inputs_embeds"
+    if "inputs_embeds" in inspect.signature(create_causal_mask).parameters
+    else "input_embeds"
+)
+
+
+def _tp_world_size() -> int:
+    if not model_parallel_is_initialized():
+        return 1
+    return get_tp_world_size()
+
+
+def _linear_output(linear: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    output = linear(x)
+    return output[0] if isinstance(output, tuple) else output
+
+
+def _make_column_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    use_tensor_parallel: bool,
+):
+    if use_tensor_parallel:
+        return ColumnParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+            gather_output=False,
+        )
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
+def _make_row_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    use_tensor_parallel: bool,
+):
+    if use_tensor_parallel:
+        return RowParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+        )
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
+def _can_use_unmasked_causal_attention(
+    attention_mask: Optional[torch.Tensor],
+    config: MistralConfig,
+    past_key_values: Optional[Cache],
+) -> bool:
+    if (
+        getattr(config, "sliding_window", None) is not None
+        or past_key_values is not None
+    ):
+        return False
+    if attention_mask is None:
+        return True
+    if attention_mask.dim() != 2:
+        return False
+    return bool(torch.all(attention_mask > 0).item())
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
+    The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to
+    (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class MistralAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: MistralConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = (
+            getattr(config, "head_dim", None)
+            or config.hidden_size // config.num_attention_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.total_num_heads = config.num_attention_heads
+        self.total_num_key_value_heads = config.num_key_value_heads
+        tp_size = _tp_world_size()
+        q_size = self.total_num_heads * self.head_dim
+        kv_size = self.total_num_key_value_heads * self.head_dim
+        self.use_tensor_parallel = (
+            tp_size > 1
+            and self.total_num_heads % tp_size == 0
+            and self.total_num_key_value_heads % tp_size == 0
+            and q_size % tp_size == 0
+            and kv_size % tp_size == 0
+        )
+        self.num_heads = (
+            self.total_num_heads // tp_size
+            if self.use_tensor_parallel
+            else self.total_num_heads
+        )
+        self.num_key_value_heads = (
+            self.total_num_key_value_heads // tp_size
+            if self.use_tensor_parallel
+            else self.total_num_key_value_heads
+        )
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.q_proj = _make_column_linear(
+            config.hidden_size,
+            q_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        self.k_proj = _make_column_linear(
+            config.hidden_size,
+            kv_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        self.v_proj = _make_column_linear(
+            config.hidden_size,
+            kv_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        self.o_proj = _make_row_linear(
+            q_size,
+            config.hidden_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        self.is_causal = True
+        self.attn = LocalAttention(
+            self.num_heads,
+            self.head_dim,
+            self.num_key_value_heads,
+            softmax_scale=self.scaling,
+            causal=True,
+            supported_attention_backends={
+                AttentionBackendEnum.FA,
+                AttentionBackendEnum.TORCH_SDPA,
+            },
+            allow_cudnn_sdp=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = (
+            _linear_output(self.q_proj, hidden_states)
+            .view(hidden_shape)
+            .transpose(1, 2)
+        )
+        key_states = (
+            _linear_output(self.k_proj, hidden_states)
+            .view(hidden_shape)
+            .transpose(1, 2)
+        )
+        value_states = (
+            _linear_output(self.v_proj, hidden_states)
+            .view(hidden_shape)
+            .transpose(1, 2)
+        )
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        attn_output = self.attn(
+            query_states.transpose(1, 2),
+            key_states.transpose(1, 2),
+            value_states.transpose(1, 2),
+            attn_mask=attention_mask,
+        )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = _linear_output(self.o_proj, attn_output)
+        return attn_output, None
+
+
+class MistralTPMLP(nn.Module):
+    def __init__(self, config: MistralConfig):
+        super().__init__()
+        tp_size = _tp_world_size()
+        use_tensor_parallel = tp_size > 1 and config.intermediate_size % tp_size == 0
+        self.gate_proj = _make_column_linear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.up_proj = _make_column_linear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.down_proj = _make_row_linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act_fn(_linear_output(self.gate_proj, x)) * _linear_output(
+            self.up_proj, x
+        )
+        return _linear_output(self.down_proj, x)
+
+
+class MistralDecoderLayer(nn.Module):
+    def __init__(self, config: MistralConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = MistralAttention(config=config, layer_idx=layer_idx)
+        self.mlp = MistralTPMLP(config)
+        self.input_layernorm = MistralRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = MistralRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[
+            tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # necessary, but kept here for BC
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class MistralModel(MistralPreTrainedModel):
+    def __init__(self, config: MistralConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                MistralDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = MistralRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+        if _can_use_unmasked_causal_attention(
+            attention_mask, self.config, past_key_values
+        ):
+            causal_mask = None
+        else:
+            mask_function = (
+                create_causal_mask
+                if getattr(self.config, "sliding_window", None) is None
+                else create_sliding_window_causal_mask
+            )
+            mask_kwargs = {
+                "config": self.config,
+                _CREATE_CAUSAL_MASK_ARG: inputs_embeds,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_mask = mask_function(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        hidden_states_pool = [] if output_hidden_states else None
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                hidden_states_pool.append(hidden_states)
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            hidden_states_pool.append(hidden_states)
+
+        return BaseModelOutputWithPast(
+            hidden_states=hidden_states_pool,
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+class Mistral3Model(nn.Module):
+    _checkpoint_conversion_mapping = {"language_model.model": "language_model"}
+
+    def __init__(self, config: Mistral3Config):
+        super().__init__()
+        self.language_model = MistralModel(config.text_config)
+        self.config = config
+
+    def get_input_embeddings(self):
+        return self.language_model.embed_tokens
+
+    def set_decoder(self, decoder):
+        self.language_model = decoder
+
+    def get_decoder(self):
+        return self.language_model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        vision_feature_layer: Optional[Union[int, list[int]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        image_sizes: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[tuple, Mistral3ModelOutputWithPast]:
+        del pixel_values, vision_feature_layer, return_dict
+        output_attentions = False
+        output_hidden_states = True
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        outputs: BaseModelOutputWithPast = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        return Mistral3ModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class Mistral3ForConditionalGeneration(nn.Module, LayerwiseOffloadableModuleMixin):
+    _checkpoint_conversion_mapping = {
+        "^language_model.model": "model.language_model",
+        "^multi_modal_projector": "model.multi_modal_projector",
+        "^language_model.lm_head": "lm_head",
+    }
+    _tied_weights_keys = ["lm_head.weight"]
+    uses_sglang_forward_context = True
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["model.language_model.layers"]
+
+    def __init__(self, config: LlavaConfig):
+        super().__init__()
+        self.model = Mistral3Model(config.arch_config)
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+
+    # Make modules available through conditional class for BC
+    @property
+    def language_model(self):
+        return self.model.language_model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_hidden_states: Optional[bool] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        image_sizes: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[tuple, Mistral3CausalLMOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        """
+        output_hidden_states = True
+
+        execution_tensor = input_ids if input_ids is not None else inputs_embeds
+        sdpa_context = (
+            sdpa_kernel(SDPBackend.CUDNN_ATTENTION)
+            if execution_tensor is not None
+            and execution_tensor.device.type == "cuda"
+            and current_platform.is_cuda()
+            else nullcontext()
+        )
+        with sdpa_context:
+            # FLUX.2 uses the text-only Mistral3 path but still expects the
+            # same local SDPA kernel choice as the official HF implementation.
+            outputs = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_hidden_states=output_hidden_states,
+                return_dict=True,
+                cache_position=cache_position,
+                image_sizes=image_sizes,
+                **kwargs,
+            )
+
+        return Mistral3CausalLMOutputWithPast(
+            hidden_states=outputs.hidden_states,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Define mapping for stacked parameters
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            name_lower = name.lower()
+            if (
+                "vision" in name_lower
+                or "multi" in name_lower
+                or "lm_head" in name_lower
+            ):
+                continue
+            final_name = name.replace("language_model.model.", "model.language_model.")
+
+            if final_name in params_dict:
+                param = params_dict[final_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(final_name)
+            else:
+                logger.warning(f"Param {name=} {final_name=} from weight is not loaded")
+
+        return loaded_params
+
+
+EntryClass = Mistral3ForConditionalGeneration

python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py

@@ -0,0 +1,487 @@
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
+
+
+class Qwen3MLP(nn.Module):
+    """Qwen3 MLP with SwiGLU activation and tensor parallelism."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. Only silu is supported."
+            )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Qwen3Attention(nn.Module):
+    """Qwen3 attention with QK-Norm and tensor parallelism.
+
+    Key difference from LLaMA: RMSNorm is applied to Q and K before attention.
+    """
+
+    def __init__(
+        self,
+        config: Qwen3TextConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 1000000.0,
+        rope_scaling: dict[str, Any] | None = None,
+        max_position_embeddings: int = 40960,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tp_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        self.head_dim = getattr(
+            config, "head_dim", self.hidden_size // self.total_num_heads
+        )
+        self.rotary_dim = self.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        # QKV projection with tensor parallelism
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        # Output projection
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # QK-Norm: Key difference from LLaMA
+        rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+        # Rotary embeddings
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.rotary_dim,
+            max_position=max_position_embeddings,
+            base=int(rope_theta),
+            rope_scaling=rope_scaling,
+            is_neox_style=True,
+        )
+
+        # Attention with FlashAttention/SageAttn support
+        self.attn = LocalAttention(
+            self.num_heads,
+            self.head_dim,
+            self.num_kv_heads,
+            softmax_scale=self.scaling,
+            causal=True,
+            supported_attention_backends=config._supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_lengths: tuple[int, ...] | None = None,
+    ) -> torch.Tensor:
+        # QKV projection
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        # Reshape for QK-norm
+        batch_size, seq_len = q.shape[0], q.shape[1]
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        # Apply QK-Norm (key difference from LLaMA)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Reshape back for rotary embeddings
+        q = q.reshape(batch_size, seq_len, -1)
+        k = k.reshape(batch_size, seq_len, -1)
+
+        # Apply rotary embeddings
+        q, k = self.rotary_emb(positions, q, k)
+
+        # Reshape for attention
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        # Attention
+        attn_output = self._masked_causal_attention(q, k, v, attention_lengths)
+        attn_output = attn_output.reshape(batch_size, seq_len, -1)
+
+        # Output projection
+        output, _ = self.o_proj(attn_output)
+        return output
+
+    def _masked_causal_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attention_lengths: tuple[int, ...] | None,
+    ) -> torch.Tensor:
+        if attention_lengths is None:
+            return self.attn(q, k, v)
+
+        seq_len = q.shape[1]
+        if all(valid_len == seq_len for valid_len in attention_lengths):
+            return self.attn(q, k, v)
+
+        outputs: list[torch.Tensor] = []
+        for batch_index, valid_len in enumerate(attention_lengths):
+            q_item = q[batch_index : batch_index + 1]
+            k_item = k[batch_index : batch_index + 1]
+            v_item = v[batch_index : batch_index + 1]
+
+            real_output = self.attn(
+                q_item[:, :valid_len],
+                k_item[:, :valid_len],
+                v_item[:, :valid_len],
+            )
+            if valid_len == seq_len:
+                outputs.append(real_output)
+                continue
+
+            pad_q = q_item[:, valid_len:].transpose(1, 2)
+            real_k = k_item[:, :valid_len].transpose(1, 2)
+            real_v = v_item[:, :valid_len].transpose(1, 2)
+            if self.num_heads != self.num_kv_heads:
+                repeat_factor = self.num_heads // self.num_kv_heads
+                real_k = real_k.repeat_interleave(repeat_factor, dim=1)
+                real_v = real_v.repeat_interleave(repeat_factor, dim=1)
+            pad_output = torch.nn.functional.scaled_dot_product_attention(
+                pad_q,
+                real_k,
+                real_v,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scaling,
+            ).transpose(1, 2)
+            outputs.append(torch.cat([real_output, pad_output], dim=1))
+
+        return torch.cat(outputs, dim=0)
+
+
+class Qwen3DecoderLayer(nn.Module):
+    """Qwen3 transformer decoder layer."""
+
+    def __init__(
+        self,
+        config: Qwen3TextConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = config.rope_parameters["rope_theta"]
+        rope_scaling = config.rope_parameters
+        max_position_embeddings = getattr(config, "max_position_embeddings", 40960)
+        attention_bias = getattr(config, "attention_bias", False)
+
+        self.self_attn = Qwen3Attention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(
+                config, "num_key_value_heads", config.num_attention_heads
+            ),
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = Qwen3MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            bias=getattr(config, "mlp_bias", False),
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        attention_lengths: tuple[int, ...] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            attention_lengths=attention_lengths,
+        )
+
+        # MLP
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class Qwen3ForCausalLM(TextEncoder):
+    """Qwen3 causal language model for text encoding in diffusion models.
+
+    Features:
+    - Tensor parallelism support
+    - FlashAttention/SageAttn/SDPA support via LocalAttention
+    - QK-Norm for better training stability
+    - FSDP sharding for CPU offload
+    """
+
+    def __init__(self, config: Qwen3TextConfig) -> None:
+        super().__init__(config)
+
+        self.config = config
+        self.quant_config = config.quant_config
+
+        # Embedding layer with tensor parallelism
+        if config.lora_config is not None:
+            max_loras = getattr(config.lora_config, "max_loras", 1)
+            lora_vocab_size = getattr(config.lora_config, "lora_extra_vocab_size", 1)
+            lora_vocab = lora_vocab_size * max_loras
+        else:
+            lora_vocab = 0
+        self.vocab_size = config.vocab_size + lora_vocab
+        self.org_vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+            org_num_embeddings=config.vocab_size,
+            quant_config=config.quant_config,
+        )
+
+        # Transformer layers
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DecoderLayer(
+                    config=config,
+                    quant_config=config.quant_config,
+                    prefix=f"{config.prefix}.layers.{i}",
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+        # Final layer norm
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+
+        residual = None
+
+        if position_ids is None:
+            position_ids = (
+                torch.arange(0, hidden_states.shape[1], device=hidden_states.device)
+                .unsqueeze(0)
+                .expand(hidden_states.shape[0], -1)
+            )
+
+        attention_lengths = None
+        if attention_mask is not None:
+            attention_lengths = tuple(
+                int(valid_len)
+                for valid_len in attention_mask.sum(dim=-1).detach().cpu().tolist()
+            )
+
+        all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
+
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                all_hidden_states += (
+                    (hidden_states,)
+                    if residual is None
+                    else (hidden_states + residual,)
+                )
+            hidden_states, residual = layer(
+                position_ids, hidden_states, residual, attention_lengths
+            )
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        # Add hidden states from the last decoder layer
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states,)
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights with support for tensor parallelism and weight remapping."""
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            # Strip 'model.' prefix from HuggingFace Qwen3 weights
+            if name.startswith("model."):
+                name = name[6:]  # len("model.") == 6
+
+            # Skip rotary embedding weights
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                continue
+
+            # Handle KV scale remapping
+            if "scale" in name:
+                kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
+                if kv_scale_name is None:
+                    continue
+                else:
+                    name = kv_scale_name
+
+            # Handle stacked params mapping (qkv_proj, gate_up_proj)
+            for (
+                param_name,
+                weight_name,
+                shard_id,
+            ) in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                # Skip loading extra bias for GPTQ models
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(name)
+
+        return loaded_params
+
+
+EntryClass = Qwen3ForCausalLM

python/sglang/multimodal_gen/runtime/models/encoders/t5.py

@@ -0,0 +1,749 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/t5/modeling_t5.py
+
+# Derived from T5 implementation posted on HuggingFace; license below:
+#
+# coding=utf-8
+# Copyright 2018 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch T5 & UMT5 model."""
+
+import math
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput, T5Config
+from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+from sglang.multimodal_gen.runtime.models.encoders.base import (
+    TextEncoder,
+    get_folding_tp_group,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+class AttentionType:
+    """
+    Attention type.
+    Use string to be compatible with `torch.compile`.
+    """
+
+    # Decoder attention between previous layer Q/K/V
+    DECODER = "decoder"
+    # Encoder attention between previous layer Q/K/V for encoder-decoder
+    ENCODER = "encoder"
+    # Encoder attention between previous layer Q/K/V
+    ENCODER_ONLY = "encoder_only"
+    # Attention between dec. Q and enc. K/V for encoder-decoder
+    ENCODER_DECODER = "encoder_decoder"
+
+
+@dataclass
+class AttentionMetadata:
+    attn_bias: torch.Tensor
+
+
+class T5DenseActDense(nn.Module):
+
+    def __init__(
+        self, config: T5Config, quant_config: QuantizationConfig | None = None
+    ):
+        super().__init__()
+        tp_group = get_folding_tp_group(config)
+        self.wi = MergedColumnParallelLinear(
+            config.d_model, [config.d_ff], bias=False, tp_group=tp_group
+        )
+        self.wo = RowParallelLinear(
+            config.d_ff,
+            config.d_model,
+            bias=False,
+            quant_config=quant_config,
+            tp_group=tp_group,
+        )
+        self.act = get_act_fn(config.dense_act_fn)
+
+    def forward(self, hidden_states) -> torch.Tensor:
+        hidden_states, _ = self.wi(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states, _ = self.wo(hidden_states)
+        return hidden_states
+
+
+class T5DenseGatedActDense(nn.Module):
+
+    def __init__(
+        self, config: T5Config, quant_config: QuantizationConfig | None = None
+    ):
+        super().__init__()
+        tp_group = get_folding_tp_group(config)
+        self.wi_0 = MergedColumnParallelLinear(
+            config.d_model,
+            [config.d_ff],
+            bias=False,
+            quant_config=quant_config,
+            tp_group=tp_group,
+        )
+        self.wi_1 = MergedColumnParallelLinear(
+            config.d_model,
+            [config.d_ff],
+            bias=False,
+            quant_config=quant_config,
+            tp_group=tp_group,
+        )
+        # Should not run in fp16 unless mixed-precision is used,
+        # see https://github.com/huggingface/transformers/issues/20287.
+        self.wo = RowParallelLinear(
+            config.d_ff,
+            config.d_model,
+            bias=False,
+            quant_config=quant_config,
+            tp_group=tp_group,
+        )
+        self.act = get_act_fn(config.dense_act_fn)
+
+    def forward(self, hidden_states) -> torch.Tensor:
+        hidden_gelu = self.act(self.wi_0(hidden_states)[0])
+        hidden_linear, _ = self.wi_1(hidden_states)
+        hidden_states = hidden_gelu * hidden_linear
+        hidden_states, _ = self.wo(hidden_states)
+        return hidden_states
+
+
+class T5LayerFF(nn.Module):
+
+    def __init__(
+        self, config: T5Config, quant_config: QuantizationConfig | None = None
+    ):
+        super().__init__()
+        if config.is_gated_act:
+            self.DenseReluDense = T5DenseGatedActDense(
+                config, quant_config=quant_config
+            )
+        else:
+            self.DenseReluDense = T5DenseActDense(config, quant_config=quant_config)
+
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(self, hidden_states) -> torch.Tensor:
+        forwarded_states = self.layer_norm(hidden_states)
+        forwarded_states = self.DenseReluDense(forwarded_states)
+        hidden_states = hidden_states + forwarded_states
+        return hidden_states
+
+
+# T5 has attn_bias and does not use softmax scaling
+class T5MultiHeadAttention(nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, q, k, v, attn_bias=None):
+        b, _, n, c = q.shape
+        attn = torch.einsum("binc,bjnc->bnij", q, k)
+        if attn_bias is not None:
+            attn += attn_bias
+
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum("bnij,bjnc->binc", attn, v)
+        x = x.reshape(b, -1, n * c)
+        return x
+
+
+class T5Attention(nn.Module):
+
+    def __init__(
+        self,
+        config: T5Config,
+        attn_type: str,
+        has_relative_attention_bias=False,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.attn_type = attn_type
+        # Cross-attention has no relative pos encoding anyway
+        self.is_decoder = attn_type == AttentionType.DECODER
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.relative_attention_num_buckets = config.relative_attention_num_buckets
+        self.relative_attention_max_distance = config.relative_attention_max_distance
+        self.d_model = config.d_model
+        self.key_value_proj_dim = config.d_kv
+        self.total_num_heads = self.total_num_kv_heads = config.num_heads
+
+        # Partition heads across multiple tensor parallel GPUs.
+        self.tp_group = get_folding_tp_group(config)
+        self.tp_world_size = get_group_size(self.tp_group)
+        assert config.num_heads % self.tp_world_size == 0
+        self.n_heads = config.num_heads // self.tp_world_size
+
+        self.inner_dim = self.n_heads * self.key_value_proj_dim
+        # No GQA in t5.
+        # self.n_kv_heads = self.n_heads
+
+        self.qkv_proj = QKVParallelLinear(
+            self.d_model,
+            self.key_value_proj_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+            tp_group=self.tp_group,
+        )
+
+        self.attn = T5MultiHeadAttention()
+
+        if self.has_relative_attention_bias:
+            self.relative_attention_bias = VocabParallelEmbedding(
+                self.relative_attention_num_buckets,
+                self.total_num_heads,
+                org_num_embeddings=self.relative_attention_num_buckets,
+                padding_size=self.relative_attention_num_buckets,
+                quant_config=quant_config,
+                tp_group=self.tp_group,
+            )
+        self.o = RowParallelLinear(
+            self.total_num_heads * self.key_value_proj_dim,
+            self.d_model,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+            tp_group=self.tp_group,
+        )
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position, bidirectional=True, num_buckets=32, max_distance=128
+    ) -> torch.Tensor:
+        """
+        Adapted from Mesh Tensorflow:
+        https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
+        Translate relative position to a bucket number for relative attention.
+        The relative position is defined as memory_position - query_position,
+        i.e. the distance in tokens from the attending position to the
+        attended-to position. If bidirectional=False, then positive relative
+        positions are invalid. We use smaller buckets for small absolute
+        relative_position and larger buckets for larger absolute
+        relative_positions. All relative positions >=max_distance map to the
+        same bucket. All relative positions <=-max_distance map to the same
+        bucket. This should allow for more graceful generalization to longer
+        sequences than the model has been trained on
+        Args:
+            relative_position: an int32 Tensor
+            bidirectional: a boolean - whether the attention is bidirectional
+            num_buckets: an integer
+            max_distance: an integer
+        Returns:
+            a Tensor with the same shape as relative_position, containing int32
+            values in the range [0, num_buckets)
+        """  # noqa: E501
+        relative_buckets = 0
+        if bidirectional:
+            num_buckets //= 2
+            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+            relative_position = torch.abs(relative_position)
+        else:
+            relative_position = -torch.min(
+                relative_position, torch.zeros_like(relative_position)
+            )
+        # now relative_position is in the range [0, inf)
+
+        # half of the buckets are for exact increments in positions
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+
+        # The other half of the buckets are for logarithmically bigger bins
+        # in positions up to max_distance
+        relative_position_if_large = max_exact + (
+            torch.log(relative_position.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+        relative_position_if_large = torch.min(
+            relative_position_if_large,
+            torch.full_like(relative_position_if_large, num_buckets - 1),
+        )
+
+        relative_buckets += torch.where(
+            is_small, relative_position, relative_position_if_large
+        )
+        return relative_buckets
+
+    def compute_bias(self, query_length, key_length, device=None) -> torch.Tensor:
+        """Compute binned relative position bias"""
+        if device is None:
+            device = self.relative_attention_bias.weight.device
+        context_position = torch.arange(query_length, dtype=torch.long, device=device)[
+            :, None
+        ]
+        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[
+            None, :
+        ]
+        # max_seq_len, nh
+        relative_position = memory_position - context_position
+        relative_position_bucket = self._relative_position_bucket(
+            relative_position,  # shape (query_length, key_length)
+            bidirectional=(not self.is_decoder),
+            num_buckets=self.relative_attention_num_buckets,
+            max_distance=self.relative_attention_max_distance,
+        )
+        values = self.relative_attention_bias(
+            relative_position_bucket
+        )  # shape (query_length, key_length, num_heads)
+        x = values.permute([2, 0, 1]).unsqueeze(
+            0
+        )  # shape (1, num_heads, query_length, key_length)
+        return x
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # (num_tokens, d_model)
+        attention_mask: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        bs, seq_len, _ = hidden_states.shape
+        num_seqs = bs
+        n, c = (
+            self.n_heads,
+            self.key_value_proj_dim,
+        )
+        qkv, _ = self.qkv_proj(hidden_states)
+        # Projection of 'own' hidden state (self-attention). No GQA here.
+        q, k, v = qkv.split(self.inner_dim, dim=-1)
+        q = q.reshape(bs, seq_len, n, c)
+        k = k.reshape(bs, seq_len, n, c)
+        v = v.reshape(bs, seq_len, n, c)
+
+        assert attn_metadata is not None
+        attn_bias = attn_metadata.attn_bias
+        # Not compatible with CP here (as all encoder-decoder models),
+        # as it assumes homogeneous batch (prefills or decodes).
+        if self.has_relative_attention_bias:
+            # Self-attention. Compute T5 relative positional encoding.
+            # The bias term is computed on longest sequence in batch. Biases
+            # for shorter sequences are slices of the longest.
+            assert self.attn_type == AttentionType.ENCODER
+            attn_bias = self.compute_bias(seq_len, seq_len).repeat(num_seqs, 1, 1, 1)
+            attn_metadata.attn_bias = attn_bias
+        else:
+            # Encoder/Decoder Self-Attention Layer, attn bias already cached.
+            assert attn_bias is not None
+
+        if attention_mask is not None:
+            attention_mask = (
+                attention_mask.view(bs, 1, 1, -1)
+                if attention_mask.ndim == 2
+                else attention_mask.unsqueeze(1)
+            )
+            mask_val = -1e4 if current_platform.is_mps() else torch.finfo(q.dtype).min
+            attn_bias.masked_fill_(attention_mask == 0, mask_val)
+
+        if self.tp_world_size > 1:
+            rank = get_group_rank(self.tp_group)
+            attn_bias = attn_bias[
+                :, rank * self.n_heads : (rank + 1) * self.n_heads, :, :
+            ]
+
+        attn_output = self.attn(q, k, v, attn_bias)
+        output, _ = self.o(attn_output)
+        return output
+
+
+class T5LayerSelfAttention(nn.Module):
+
+    def __init__(
+        self,
+        config,
+        has_relative_attention_bias=False,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.SelfAttention = T5Attention(
+            config,
+            AttentionType.DECODER if "decoder" in prefix else AttentionType.ENCODER,
+            has_relative_attention_bias=has_relative_attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.SelfAttention",
+        )
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        normed_hidden_states = self.layer_norm(hidden_states)
+
+        attention_output = self.SelfAttention(
+            hidden_states=normed_hidden_states,
+            attention_mask=attention_mask,
+            attn_metadata=attn_metadata,
+        )
+
+        hidden_states = hidden_states + attention_output
+
+        return hidden_states
+
+
+class T5LayerCrossAttention(nn.Module):
+
+    def __init__(
+        self, config, quant_config: QuantizationConfig | None = None, prefix: str = ""
+    ):
+        super().__init__()
+        self.EncDecAttention = T5Attention(
+            config,
+            AttentionType.ENCODER_DECODER,
+            has_relative_attention_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.EncDecAttention",
+        )
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        normed_hidden_states = self.layer_norm(hidden_states)
+        attention_output = self.EncDecAttention(
+            hidden_states=normed_hidden_states,
+            attn_metadata=attn_metadata,
+        )
+        hidden_states = hidden_states + attention_output
+        return hidden_states
+
+
+class T5Block(nn.Module):
+
+    def __init__(
+        self,
+        config: T5Config,
+        is_decoder: bool,
+        has_relative_attention_bias=False,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.is_decoder = is_decoder
+        self.layer = nn.ModuleList()
+        self.layer.append(
+            T5LayerSelfAttention(
+                config,
+                has_relative_attention_bias=has_relative_attention_bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+            )
+        )
+
+        if self.is_decoder:
+            self.layer.append(
+                T5LayerCrossAttention(
+                    config, quant_config=quant_config, prefix=f"{prefix}.cross_attn"
+                )
+            )
+
+        self.layer.append(T5LayerFF(config, quant_config=quant_config))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                hidden_states.shape[:2], device=hidden_states.device
+            )
+
+        hidden_states = self.layer[0](
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            attn_metadata=attn_metadata,
+        )
+
+        if self.is_decoder:
+            hidden_states = self.layer[1](
+                hidden_states=hidden_states, attn_metadata=attn_metadata
+            )
+
+        # Apply Feed Forward layer
+        hidden_states = self.layer[-1](hidden_states)
+
+        return hidden_states
+
+
+class T5Stack(nn.Module):
+
+    def __init__(
+        self,
+        config: T5Config,
+        is_decoder: bool,
+        n_layers: int,
+        embed_tokens=None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        is_umt5: bool = False,
+    ):
+        super().__init__()
+        self.embed_tokens = embed_tokens
+        self.is_umt5 = is_umt5
+        if is_umt5:
+            self.block = nn.ModuleList(
+                [
+                    T5Block(
+                        config,
+                        is_decoder=is_decoder,
+                        has_relative_attention_bias=True,
+                        quant_config=quant_config,
+                        prefix=f"{prefix}.blocks.{i}",
+                    )
+                    for i in range(n_layers)
+                ]
+            )
+        else:
+            # Only the first block has relative positional encoding.
+            self.block = nn.ModuleList(
+                [
+                    T5Block(
+                        config,
+                        is_decoder=is_decoder,
+                        has_relative_attention_bias=i == 0,
+                        quant_config=quant_config,
+                        prefix=f"{prefix}.blocks.{i}",
+                    )
+                    for i in range(n_layers)
+                ]
+            )
+        self.final_layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+
+        for idx, block in enumerate(self.block):
+            hidden_states = block(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                attn_metadata=attn_metadata,
+            )
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        return hidden_states
+
+
+class T5EncoderModel(TextEncoder):
+
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__(config)
+
+        quant_config = None
+        tp_group = get_folding_tp_group(config)
+        self.shared = VocabParallelEmbedding(
+            config.vocab_size,
+            config.d_model,
+            org_num_embeddings=config.vocab_size,
+            tp_group=tp_group,
+        )
+
+        self.encoder = T5Stack(
+            config,
+            False,
+            config.num_layers,
+            self.shared,
+            quant_config=quant_config,
+            prefix=f"{prefix}.encoder",
+            is_umt5=False,
+        )
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        attn_metadata = AttentionMetadata(None)
+        hidden_states = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            attn_metadata=attn_metadata,
+        )
+
+        return BaseEncoderOutput(last_hidden_state=hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q", "q"),
+            (".qkv_proj", ".k", "k"),
+            (".qkv_proj", ".v", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            loaded = False
+            if "decoder" in name or "lm_head" in name:
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded = True
+                break
+            if not loaded:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class UMT5EncoderModel(TextEncoder):
+
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__(config)
+
+        quant_config = None
+        tp_group = get_folding_tp_group(config)
+        self.shared = VocabParallelEmbedding(
+            config.vocab_size,
+            config.d_model,
+            org_num_embeddings=config.vocab_size,
+            tp_group=tp_group,
+        )
+
+        self.encoder = T5Stack(
+            config,
+            False,
+            config.num_layers,
+            self.shared,
+            quant_config=quant_config,
+            prefix=f"{prefix}.encoder",
+            is_umt5=True,
+        )
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        attn_metadata = AttentionMetadata(None)
+        hidden_states = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            attn_metadata=attn_metadata,
+        )
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            attention_mask=attention_mask,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            loaded = False
+            if "decoder" in name or "lm_head" in name:
+                continue
+            for (
+                param_name,
+                weight_name,
+                shard_id,
+            ) in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded = True
+                break
+            if not loaded:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+EntryClass = [T5EncoderModel, UMT5EncoderModel]

python/sglang/multimodal_gen/runtime/models/encoders/vision.py

@@ -0,0 +1,96 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/vision.py
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+import torch
+from transformers import PretrainedConfig
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_C = TypeVar("_C", bound=PretrainedConfig)
+
+
+class VisionEncoderInfo(ABC, Generic[_C]):
+
+    def __init__(self, vision_config: _C) -> None:
+        super().__init__()
+
+        self.vision_config = vision_config
+
+    @abstractmethod
+    def get_num_image_tokens(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+    ) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_max_image_tokens(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_image_size(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_patch_size(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_patch_grid_length(self) -> int:
+        raise NotImplementedError
+
+
+def resolve_visual_encoder_outputs(
+    encoder_outputs: torch.Tensor | list[torch.Tensor],
+    feature_sample_layers: list[int] | None,
+    post_layer_norm: torch.nn.LayerNorm | None,
+    max_possible_layers: int,
+) -> torch.Tensor:
+    """Given the outputs a visual encoder module that may correspond to the
+    output of the last layer, or a list of hidden states to be stacked,
+    handle post normalization and resolve it into a single output tensor.
+
+    Args:
+        encoder_outputs: Output of encoder's last layer or all hidden states.
+        feature_sample_layers: Optional layer indices to grab from the encoder
+            outputs; if provided, encoder outputs must be a list.
+        post_layer_norm: Post norm to apply to the output of the encoder.
+        max_possible_layers: Total layers in the fully loaded visual encoder.
+
+    """
+    if feature_sample_layers is None:
+        if post_layer_norm is not None:
+            return post_layer_norm(encoder_outputs)
+        return encoder_outputs
+
+    # Get the hidden states corresponding to the layer indices.
+    # Negative values are relative to the full visual encoder,
+    # so offset them depending on how many layers were loaded.
+    # NOTE: this assumes that encoder_outputs is a list containing
+    # the inputs to the visual encoder, followed by the hidden states
+    # of each layer.
+    num_loaded_layers = len(encoder_outputs) - 1
+    offset = max_possible_layers - num_loaded_layers
+    hs_pool = [
+        (
+            encoder_outputs[layer_idx]
+            if layer_idx >= 0
+            else encoder_outputs[layer_idx + offset]
+        )
+        for layer_idx in feature_sample_layers
+    ]
+
+    # Apply post-norm on the final hidden state if we are using it
+    uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1)
+    if post_layer_norm is not None and uses_last_layer:
+        hs_pool[-1] = post_layer_norm(encoder_outputs)
+    return torch.cat(hs_pool, dim=-1)

python/sglang/multimodal_gen/runtime/models/parameter.py

@@ -0,0 +1,424 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/parameter.py
+
+from collections.abc import Callable
+from fractions import Fraction
+from typing import Any
+
+import torch
+from torch.nn import Parameter
+
+from sglang.multimodal_gen.runtime.distributed import get_tp_rank
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.weight_attrs import make_synced_weight_loader
+
+logger = init_logger(__name__)
+
+
+class BasevLLMParameter(Parameter):
+    """
+    Base parameter for vLLM linear layers. Extends the torch.nn.parameter
+    by taking in a linear weight loader. Will copy the loaded weight
+    into the parameter when the provided weight loader is called.
+    """
+
+    def __new__(cls, data: torch.Tensor, **kwargs):
+
+        return super().__new__(cls, data=data, requires_grad=False)
+
+    def __init__(self, data: torch.Tensor, weight_loader: Callable):
+        """
+        Initialize the BasevLLMParameter
+
+        :param data: torch tensor with the parameter data
+        :param weight_loader: weight loader callable
+
+        :returns: a torch.nn.parameter
+        """
+
+        # During weight loading, we often do something like:
+        # narrowed_tensor = param.data.narrow(0, offset, len)
+        # narrowed_tensor.copy_(real_weight)
+        # expecting narrowed_tensor and param.data to share the same storage.
+        # However, on TPUs, narrowed_tensor will lazily propagate to the base
+        # tensor, which is param.data, leading to the redundant memory usage.
+        # This sometimes causes OOM errors during model loading. To avoid this,
+        # we sync the param tensor after its weight loader is called.
+        from sglang.multimodal_gen.runtime.platforms import current_platform
+
+        if current_platform.is_tpu():
+            weight_loader = make_synced_weight_loader(weight_loader)
+
+        self._weight_loader = weight_loader
+
+    @property
+    def weight_loader(self):
+        return self._weight_loader
+
+    def _is_1d_and_scalar(self, loaded_weight: torch.Tensor):
+        cond1 = self.data.ndim == 1 and self.data.numel() == 1
+        cond2 = loaded_weight.ndim == 0 and loaded_weight.numel() == 1
+        return cond1 and cond2
+
+    def _assert_and_load(self, loaded_weight: torch.Tensor) -> None:
+        assert self.data.shape == loaded_weight.shape or self._is_1d_and_scalar(
+            loaded_weight
+        )
+        self.data.copy_(loaded_weight)
+
+    def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None:
+        self._assert_and_load(loaded_weight)
+
+    def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None:
+        self._assert_and_load(loaded_weight)
+
+    def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None:
+        self._assert_and_load(loaded_weight)
+
+    def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None:
+        self._assert_and_load(loaded_weight)
+
+
+class _ColumnvLLMParameter(BasevLLMParameter):
+    """
+    Private class defining weight loading functionality
+    (load_merged_column_weight, load_qkv_weight)
+    for parameters being loaded into linear layers with column
+    parallelism. This includes QKV and MLP layers which are
+    not already fused on disk. Requires an output dimension
+    to be defined. Called within the weight loader of
+    each of the column parallel linear layers.
+    """
+
+    def __init__(self, output_dim: int, **kwargs):
+        self._output_dim = output_dim
+        super().__init__(**kwargs)
+
+    @property
+    def output_dim(self):
+        return self._output_dim
+
+    def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None:
+        tp_rank = get_tp_rank()
+        shard_size = self.data.shape[self.output_dim]
+        loaded_weight = loaded_weight.narrow(
+            self.output_dim, tp_rank * shard_size, shard_size
+        )
+        assert self.data.shape == loaded_weight.shape
+        self.data.copy_(loaded_weight)
+
+    def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None:
+
+        shard_offset = kwargs.get("shard_offset")
+        shard_size = kwargs.get("shard_size")
+        if shard_offset is None or shard_size is None:
+            raise ValueError("shard_offset and shard_size must be provided")
+        if (
+            isinstance(self, PackedColumnParameter | PackedvLLMParameter)
+            and self.packed_dim == self.output_dim
+        ):
+            shard_size, shard_offset = self.adjust_shard_indexes_for_packing(
+                shard_offset=shard_offset, shard_size=shard_size
+            )
+
+        param_data = self.data
+
+        tp_rank = get_tp_rank()
+        param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
+        loaded_weight = loaded_weight.narrow(
+            self.output_dim, tp_rank * shard_size, shard_size
+        )
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+    def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None:
+
+        shard_offset = kwargs.get("shard_offset")
+        shard_size = kwargs.get("shard_size")
+        shard_id = kwargs.get("shard_id")
+        num_heads = kwargs.get("num_heads")
+
+        assert shard_offset is not None
+        assert shard_size is not None
+        assert shard_id is not None
+        assert num_heads is not None
+
+        if (
+            isinstance(self, PackedColumnParameter | PackedvLLMParameter)
+            and self.output_dim == self.packed_dim
+        ):
+            shard_size, shard_offset = self.adjust_shard_indexes_for_packing(
+                shard_offset=shard_offset, shard_size=shard_size
+            )
+
+        param_data = self.data
+        tp_rank = get_tp_rank()
+        shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads
+        param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
+        loaded_weight = loaded_weight.narrow(
+            self.output_dim, shard_id * shard_size, shard_size
+        )
+
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
+class RowvLLMParameter(BasevLLMParameter):
+    """
+    Parameter class defining weight_loading functionality
+    (load_row_parallel_weight) for parameters being loaded
+    into linear layers with row parallel functionality.
+    Requires an input_dim to be defined.
+    """
+
+    def __init__(self, input_dim: int, **kwargs):
+        self._input_dim = input_dim
+        super().__init__(**kwargs)
+
+    @property
+    def input_dim(self):
+        return self._input_dim
+
+    def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None:
+        tp_rank = get_tp_rank()
+        shard_size = self.data.shape[self.input_dim]
+        loaded_weight = loaded_weight.narrow(
+            self.input_dim, tp_rank * shard_size, shard_size
+        )
+
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+
+        assert self.data.shape == loaded_weight.shape
+        self.data.copy_(loaded_weight)
+
+
+class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter):
+    """
+    Parameter class for linear layer weights. Uses both column and
+    row parallelism.
+    """
+
+    pass
+
+
+class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
+    """
+    Parameter class for weight scales loaded for weights with
+    grouped quantization. Uses both column and row parallelism.
+    """
+
+    pass
+
+
+class ChannelQuantScaleParameter(_ColumnvLLMParameter):
+    """
+    Parameter class for weight scales loaded for weights with
+    channel-wise quantization. Equivalent to _ColumnvLLMParameter.
+    """
+
+    pass
+
+
+class PerTensorScaleParameter(BasevLLMParameter):
+    """
+    Parameter class for scales where the number of scales is
+    equivalent to the number of logical matrices in fused linear
+    layers (e.g. for QKV, there are 3 scales loaded from disk).
+    This is relevant to weights with per-tensor quantization.
+    Adds functionality to map the scalers to a shard during
+    weight loading.
+
+    Note: additional parameter manipulation may be handled
+    for each quantization config specifically, within
+    process_weights_after_loading
+    """
+
+    def __init__(self, **kwargs):
+        self.qkv_idxs = {"q": 0, "k": 1, "v": 2}
+        super().__init__(**kwargs)
+
+    def _shard_id_as_int(self, shard_id: str | int) -> int:
+        if isinstance(shard_id, int):
+            return shard_id
+
+        # if not int, assume shard_id for qkv
+        # map to int and return
+        assert isinstance(shard_id, str)
+        assert shard_id in self.qkv_idxs
+        return self.qkv_idxs[shard_id]
+
+    # For row parallel layers, no sharding needed
+    # load weight into parameter as is
+    def load_row_parallel_weight(self, *args, **kwargs) -> None:
+        super().load_row_parallel_weight(*args, **kwargs)
+
+    def load_merged_column_weight(self, *args, **kwargs) -> None:
+        self._load_into_shard_id(*args, **kwargs)
+
+    def load_qkv_weight(self, *args, **kwargs) -> None:
+        self._load_into_shard_id(*args, **kwargs)
+
+    def load_column_parallel_weight(self, *args, **kwargs) -> None:
+        super().load_row_parallel_weight(*args, **kwargs)
+
+    def _load_into_shard_id(
+        self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs
+    ):
+        """
+        Slice the parameter data based on the shard id for
+        loading.
+        """
+
+        param_data = self.data
+        shard_id = self._shard_id_as_int(shard_id)
+
+        # AutoFP8 scales do not have a shape
+        # compressed-tensors scales do have a shape
+        if len(loaded_weight.shape) != 0:
+            assert loaded_weight.shape[0] == 1
+            loaded_weight = loaded_weight[0]
+
+        param_data = param_data[shard_id]
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
+class PackedColumnParameter(_ColumnvLLMParameter):
+    """
+    Parameter for model parameters which are packed on disk
+    and support column parallelism only. See PackedvLLMParameter
+    for more details on the packed properties.
+    """
+
+    def __init__(self, packed_factor: int | Fraction, packed_dim: int, **kwargs):
+        self._packed_factor = packed_factor
+        self._packed_dim = packed_dim
+        super().__init__(**kwargs)
+
+    @property
+    def packed_dim(self):
+        return self._packed_dim
+
+    @property
+    def packed_factor(self):
+        return self._packed_factor
+
+    def adjust_shard_indexes_for_packing(
+        self, shard_size, shard_offset
+    ) -> tuple[Any, Any]:
+        return _adjust_shard_indexes_for_packing(
+            shard_size=shard_size,
+            shard_offset=shard_offset,
+            packed_factor=self.packed_factor,
+        )
+
+
+class PackedvLLMParameter(ModelWeightParameter):
+    """
+    Parameter for model weights which are packed on disk.
+    Example: GPTQ Marlin weights are int4 or int8, packed into int32.
+    Extends the ModelWeightParameter to take in the
+    packed factor, the packed dimension, and optionally, marlin
+    tile size for marlin kernels. Adjusts the shard_size and
+    shard_offset for fused linear layers model weight loading
+    by accounting for packing and optionally, marlin tile size.
+    """
+
+    def __init__(self, packed_factor: int | Fraction, packed_dim: int, **kwargs):
+        self._packed_factor = packed_factor
+        self._packed_dim = packed_dim
+        super().__init__(**kwargs)
+
+    @property
+    def packed_dim(self):
+        return self._packed_dim
+
+    @property
+    def packed_factor(self):
+        return self._packed_factor
+
+    def adjust_shard_indexes_for_packing(self, shard_size, shard_offset):
+        return _adjust_shard_indexes_for_packing(
+            shard_size=shard_size,
+            shard_offset=shard_offset,
+            packed_factor=self.packed_factor,
+        )
+
+
+class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
+    """
+    Parameter class for weight scales loaded for weights with
+    block-wise quantization. Uses both column and row parallelism.
+    """
+
+    pass
+
+
+def permute_param_layout_(
+    param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs
+) -> BasevLLMParameter:
+    """
+    Permute a parameter's layout to the specified input and output dimensions,
+    useful for forcing the parameter into a known layout, for example, if I need
+    a packed (quantized) weight matrix to be in the layout
+        {input_dim = 0, output_dim = 1, packed_dim = 0}
+    then I can call:
+        permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+    to ensure x is in the correct layout (permuting it to the correct layout if
+    required, asserting if it cannot get it to the correct layout)
+    """
+
+    curr_input_dim = getattr(param, "input_dim", None)
+    curr_output_dim = getattr(param, "output_dim", None)
+
+    if curr_input_dim is None or curr_output_dim is None:
+        assert param.data.dim() == 2, (
+            "permute_param_layout_ only supports 2D parameters when either "
+            "input_dim or output_dim is not set"
+        )
+
+    # if one of the dimensions is not set, set it to the opposite of the other
+    #  we can only do this since we asserted the parameter is 2D above
+    if curr_input_dim is None:
+        assert curr_output_dim is not None, "either input or output dim must be set"
+        curr_input_dim = (curr_output_dim + 1) % 2
+    if curr_output_dim is None:
+        assert curr_input_dim is not None, "either input or output dim must be set"
+        curr_output_dim = (curr_input_dim + 1) % 2
+
+    # create permutation from the current layout to the layout with
+    # self.input_dim at input_dim and self.output_dim at output_dim preserving
+    # other dimensions
+    perm = [
+        i for i in range(param.data.dim()) if i not in [curr_input_dim, curr_output_dim]
+    ]
+    perm.insert(input_dim, curr_input_dim)
+    perm.insert(output_dim, curr_output_dim)
+
+    if "packed_dim" in kwargs:
+        assert (
+            hasattr(param, "packed_dim")
+            and param.packed_dim == perm[kwargs["packed_dim"]]
+        ), "permute_param_layout_ currently doesn't support repacking"
+
+    param.data = param.data.permute(*perm)
+    if hasattr(param, "_input_dim"):
+        param._input_dim = input_dim
+    if hasattr(param, "_output_dim"):
+        param._output_dim = output_dim
+    if "packed_dim" in kwargs and hasattr(param, "_packed_dim"):
+        param._packed_dim = kwargs["packed_dim"]
+
+    return param
+
+
+def _adjust_shard_indexes_for_packing(
+    shard_size, shard_offset, packed_factor
+) -> tuple[Any, Any]:
+    shard_size = shard_size // packed_factor
+    shard_offset = shard_offset // packed_factor
+    return shard_size, shard_offset

python/sglang/multimodal_gen/runtime/models/registry.py

@@ -0,0 +1,422 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/registry.py
+
+import ast
+import importlib
+import os
+import pickle
+import subprocess
+import sys
+import tempfile
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Set
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import NoReturn, TypeVar, cast
+
+import cloudpickle
+from torch import nn
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+MODELS_PATH = os.path.dirname(__file__)
+COMPONENT_DIRS = [
+    d
+    for d in os.listdir(MODELS_PATH)
+    if os.path.isdir(os.path.join(MODELS_PATH, d))
+    and not d.startswith("__")
+    and not d.startswith(".")
+]
+
+_IMAGE_ENCODER_MODELS: dict[str, tuple] = {
+    # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"),
+    "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"),
+}
+
+# Global alias mapping: external_path -> canonical_class_name
+_ALIAS_TO_MODEL: dict[str, str] = {}
+
+
+def _parse_aliases_from_ast(value_node: ast.expr) -> list[str]:
+    """Parse _aliases list from AST node."""
+    aliases = []
+    if isinstance(value_node, (ast.List, ast.Tuple)):
+        for elt in value_node.elts:
+            if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
+                aliases.append(elt.value)
+    return aliases
+
+
+@lru_cache(maxsize=None)
+def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
+    discovered_models = dict(_IMAGE_ENCODER_MODELS)
+
+    # Collect class definitions with their _aliases
+    class_aliases: dict[str, list[str]] = {}
+
+    for component in COMPONENT_DIRS:
+        component_path = os.path.join(MODELS_PATH, component)
+        for filename in os.listdir(component_path):
+            if not filename.endswith(".py"):
+                continue
+
+            mod_relname = filename[:-3]
+            filepath = os.path.join(component_path, filename)
+            try:
+                with open(filepath, "r", encoding="utf-8") as f:
+                    source = f.read()
+                tree = ast.parse(source, filename=filename)
+
+                entry_class_node = None
+                first_class_def = None
+
+                # Collect all class definitions and their _aliases
+                file_class_aliases: dict[str, list[str]] = {}
+                for node in ast.walk(tree):
+                    if isinstance(node, ast.ClassDef):
+                        if first_class_def is None:
+                            first_class_def = node
+                        # Look for _aliases in the class body
+                        for class_body_node in node.body:
+                            if isinstance(class_body_node, ast.Assign):
+                                for target in class_body_node.targets:
+                                    if (
+                                        isinstance(target, ast.Name)
+                                        and target.id == "_aliases"
+                                    ):
+                                        aliases = _parse_aliases_from_ast(
+                                            class_body_node.value
+                                        )
+                                        if aliases:
+                                            file_class_aliases[node.name] = aliases
+                    if isinstance(node, ast.Assign):
+                        for target in node.targets:
+                            if (
+                                isinstance(target, ast.Name)
+                                and target.id == "EntryClass"
+                            ):
+                                entry_class_node = node
+                                break
+
+                if entry_class_node and first_class_def:
+                    model_cls_name_list = []
+                    value_node = entry_class_node.value
+
+                    # EntryClass = ClassName
+                    if isinstance(value_node, ast.Name):
+                        model_cls_name_list.append(value_node.id)
+                    # EntryClass = ["...", ClassName, ...]
+                    elif isinstance(value_node, (ast.List, ast.Tuple)):
+                        for elt in value_node.elts:
+                            if isinstance(elt, ast.Constant):
+                                model_cls_name_list.append(elt.value)
+                            elif isinstance(elt, ast.Name):
+                                model_cls_name_list.append(elt.id)
+
+                    if model_cls_name_list:
+                        for model_cls_str in model_cls_name_list:
+                            if model_cls_str in discovered_models:
+                                logger.warning(
+                                    f"Duplicate architecture found: {model_cls_str}. It will be overwritten."
+                                )
+                            model_arch = model_cls_str
+                            discovered_models[model_arch] = (
+                                component,
+                                mod_relname,
+                                model_cls_str,
+                            )
+                            # Collect aliases for this class
+                            if model_cls_str in file_class_aliases:
+                                class_aliases[model_cls_str] = file_class_aliases[
+                                    model_cls_str
+                                ]
+
+            except Exception as e:
+                logger.warning(f"Could not parse {filepath} to find models: {e}")
+
+    # Build alias -> canonical class name mapping
+    for class_name, aliases in class_aliases.items():
+        for alias in aliases:
+            if alias in _ALIAS_TO_MODEL:
+                logger.warning(
+                    f"Alias '{alias}' already registered for '{_ALIAS_TO_MODEL[alias]}', "
+                    f"will be overwritten by '{class_name}'"
+                )
+            _ALIAS_TO_MODEL[alias] = class_name
+
+    return discovered_models
+
+
+_SGLANG_DIFFUSION_MODELS = _discover_and_register_models()
+
+_SUBPROCESS_COMMAND = [
+    sys.executable,
+    "-m",
+    "sglang.multimodal_gen.runtime.models.dits.registry",
+]
+
+_T = TypeVar("_T")
+
+
+@dataclass(frozen=True)
+class _ModelInfo:
+    architecture: str
+
+    @staticmethod
+    def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
+        return _ModelInfo(
+            architecture=model.__name__,
+        )
+
+
+class _BaseRegisteredModel(ABC):
+
+    @abstractmethod
+    def inspect_model_cls(self) -> _ModelInfo:
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_model_cls(self) -> type[nn.Module]:
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class _RegisteredModel(_BaseRegisteredModel):
+    """
+    Represents a model that has already been imported in the main process.
+    """
+
+    interfaces: _ModelInfo
+    model_cls: type[nn.Module]
+
+    @staticmethod
+    def from_model_cls(model_cls: type[nn.Module]):
+        return _RegisteredModel(
+            interfaces=_ModelInfo.from_model_cls(model_cls),
+            model_cls=model_cls,
+        )
+
+    def inspect_model_cls(self) -> _ModelInfo:
+        return self.interfaces
+
+    def load_model_cls(self) -> type[nn.Module]:
+        return self.model_cls
+
+
+def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
+    # NOTE: We use a temporary directory instead of a temporary file to avoid
+    # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
+    with tempfile.TemporaryDirectory() as tempdir:
+        output_filepath = os.path.join(tempdir, "registry_output.tmp")
+
+        # `cloudpickle` allows pickling lambda functions directly
+        input_bytes = cloudpickle.dumps((fn, output_filepath))
+
+        # cannot use `sys.executable __file__` here because the script
+        # contains relative imports
+        returned = subprocess.run(
+            _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True
+        )
+
+        # check if the subprocess is successful
+        try:
+            returned.check_returncode()
+        except Exception as e:
+            # wrap raised exception to provide more information
+            raise RuntimeError(
+                f"Error raised in subprocess:\n" f"{returned.stderr.decode()}"
+            ) from e
+
+        with open(output_filepath, "rb") as f:
+            return cast(_T, pickle.load(f))
+
+
+@dataclass(frozen=True)
+class _LazyRegisteredModel(_BaseRegisteredModel):
+    """
+    Represents a model that has not been imported in the main process.
+    """
+
+    module_name: str
+    component_name: str
+    class_name: str
+
+    # Performed in another process to avoid initializing CUDA
+    def inspect_model_cls(self) -> _ModelInfo:
+        return _run_in_subprocess(
+            lambda: _ModelInfo.from_model_cls(self.load_model_cls())
+        )
+
+    def load_model_cls(self) -> type[nn.Module]:
+        mod = importlib.import_module(self.module_name)
+        return cast(type[nn.Module], getattr(mod, self.class_name))
+
+
+@lru_cache(maxsize=128)
+def _try_load_model_cls(
+    model_arch: str,
+    model: _BaseRegisteredModel,
+) -> type[nn.Module] | None:
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    current_platform.verify_model_arch(model_arch)
+    try:
+        return model.load_model_cls()
+    except Exception:
+        logger.exception("Ignore import error when loading '%s'", model_arch)
+        return None
+
+
+@lru_cache(maxsize=128)
+def _try_inspect_model_cls(
+    model_arch: str,
+    model: _BaseRegisteredModel,
+) -> _ModelInfo | None:
+    try:
+        return model.inspect_model_cls()
+    except Exception:
+        logger.exception("Error in inspecting model architecture '%s'", model_arch)
+        return None
+
+
+@dataclass
+class _ModelRegistry:
+    # Keyed by model_arch
+    registered_models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
+
+    def get_supported_archs(self) -> Set[str]:
+        return self.registered_models.keys()
+
+    def resolve_by_alias(self, alias: str) -> type[nn.Module] | None:
+        """Resolve a model class by its alias (external module path)."""
+        if alias in _ALIAS_TO_MODEL:
+            canonical_name = _ALIAS_TO_MODEL[alias]
+            return self._try_load_model_cls(canonical_name)
+        return None
+
+    def register_model(
+        self,
+        model_arch: str,
+        model_cls: type[nn.Module] | str,
+    ) -> None:
+        """
+        Register an external model to be used in vLLM.
+
+        :code:`model_cls` can be either:
+
+        - A :class:`torch.nn.Module` class directly referencing the model.
+        - A string in the format :code:`<module>:<class>` which can be used to
+          lazily import the model. This is useful to avoid initializing CUDA
+          when importing the model and thus the related error
+          :code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
+        """
+        if model_arch in self.registered_models:
+            logger.warning(
+                "Model architecture %s is already registered, and will be "
+                "overwritten by the new model class %s.",
+                model_arch,
+                model_cls,
+            )
+
+        if isinstance(model_cls, str):
+            split_str = model_cls.split(":")
+            if len(split_str) != 2:
+                msg = "Expected a string in the format `<module>:<class>`"
+                raise ValueError(msg)
+
+            model = _LazyRegisteredModel(*split_str)
+        else:
+            model = _RegisteredModel.from_model_cls(model_cls)
+
+        self.registered_models[model_arch] = model
+
+    def _raise_for_unsupported(self, architectures: list[str]) -> NoReturn:
+        all_supported_archs = self.get_supported_archs()
+
+        if any(arch in all_supported_archs for arch in architectures):
+            raise ValueError(
+                f"Model architectures {architectures} failed "
+                "to be inspected. Please check the logs for more details."
+            )
+
+        raise ValueError(
+            f"Model architectures {architectures} are not supported for now. "
+            f"Supported architectures: {all_supported_archs}"
+        )
+
+    def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None:
+        if model_arch not in self.registered_models:
+            return None
+
+        return _try_load_model_cls(model_arch, self.registered_models[model_arch])
+
+    def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None:
+        if model_arch not in self.registered_models:
+            return None
+
+        return _try_inspect_model_cls(model_arch, self.registered_models[model_arch])
+
+    def _normalize_archs(
+        self,
+        architectures: str | list[str],
+    ) -> list[str]:
+        if isinstance(architectures, str):
+            architectures = [architectures]
+        if not architectures:
+            logger.warning("No model architectures are specified")
+
+        normalized_arch = []
+        for arch in architectures:
+            if arch not in self.registered_models:
+                registered_models = list(self.registered_models.keys())
+                raise Exception(
+                    f"Unsupported model architecture: {arch}. Registered architectures: {registered_models}"
+                )
+            normalized_arch.append(arch)
+        return normalized_arch
+
+    def inspect_model_cls(
+        self,
+        architectures: str | list[str],
+    ) -> tuple[_ModelInfo, str]:
+        architectures = self._normalize_archs(architectures)
+
+        for arch in architectures:
+            model_info = self._try_inspect_model_cls(arch)
+            if model_info is not None:
+                return (model_info, arch)
+
+        return self._raise_for_unsupported(architectures)
+
+    def resolve_model_cls(
+        self,
+        architectures: str | list[str],
+    ) -> tuple[type[nn.Module], str]:
+        architectures = self._normalize_archs(architectures)
+
+        for arch in architectures:
+            model_cls = self._try_load_model_cls(arch)
+            if model_cls is not None:
+                return (model_cls, arch)
+
+        return self._raise_for_unsupported(architectures)
+
+
+ModelRegistry = _ModelRegistry(
+    {
+        model_arch: _LazyRegisteredModel(
+            module_name=f"sglang.multimodal_gen.runtime.models.{component_name}.{mod_relname}",
+            component_name=component_name,
+            class_name=cls_name,
+        )
+        for model_arch, (
+            component_name,
+            mod_relname,
+            cls_name,
+        ) in _SGLANG_DIFFUSION_MODELS.items()
+    }
+)

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

@@ -0,0 +1,37 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+import torch
+
+
+class BaseScheduler(ABC):
+    timesteps: torch.Tensor
+    order: int
+    num_train_timesteps: int
+
+    def __init__(self, *args, **kwargs) -> None:
+        # Check if subclass has defined all required properties
+        required_attributes = ["timesteps", "order", "num_train_timesteps"]
+
+        for attr in required_attributes:
+            if not hasattr(self, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseScheduler must define '{attr}' property"
+                )
+
+    @abstractmethod
+    def set_shift(self, shift: float) -> None:
+        pass
+
+    @abstractmethod
+    def set_timesteps(self, *args, **kwargs) -> None:
+        pass
+
+    @abstractmethod
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        pass

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

@@ -0,0 +1,531 @@
+# Copied and adapted from: https://github.com/OpenMOSS/MOVA/tree/main/mova/diffusion/schedulers/flow_match.py and flow_match_pair.py
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+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)
+            if scale_factor != 0:
+                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
+        # Initialize train_timesteps on first set.
+        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):
+    """Pairing scheduler built on FlowMatchScheduler.
+
+    Provides a convenient pairing interface for timesteps or sigmas.
+
+    Attributes:
+        pair_timesteps: Cached timestep pairs of shape [num_timesteps, 2].
+        pair_sigmas: Cached sigma pairs of shape [num_timesteps, 2].
+    """
+
+    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):
+        """Set a postprocess function to customize pairs after construction.
+
+        Args:
+            fn: Callable with signature fn(pairs: torch.Tensor) -> torch.Tensor.
+                The returned tensor must have the same shape as input pairs.
+
+        Raises:
+            TypeError: If fn is not callable or None.
+            RuntimeError: If scheduler is not initialized.
+        """
+        if fn is not None and not callable(fn):
+            raise TypeError("pair_postprocess must be callable or 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("Scheduler not initialized; call set_timesteps() first")
+        self._refresh_pair_cache()
+
+    def set_pair_postprocess_by_name(self, name: str | None, **kwargs):
+        """Configure a postprocess function by name.
+
+        Supported names:
+            - None/"none"/"off"/"false"/"no": disable
+            - "quadratic_perp_bulge_swap": x2=x+d, y2=x-d, where d=4*amp*s*(1-s), s=t/T
+            - "v2a_sequential": assume pairs are (t,t); sample half sequence from column 0
+              with stride 2, then let column 0 follow this sequence first, followed by column 1
+            - "a2v_sequential": same as above, but column 1 first then column 0
+            - "dual_sigma_shift": use only timestep count; rebuild two columns independently using
+              FlowMatchScheduler sigma transform logic; configurable visual_shift/audio_shift
+
+        Args:
+            name: Postprocess name or None to disable.
+            **kwargs: Extra parameters for the named postprocess. For example:
+                - amp: Float amplitude, default 150.0.
+
+        Raises:
+            ValueError: If name is unknown.
+        """
+
+        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 must be a torch.Tensor of shape [N, 2]")
+                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 must be a torch.Tensor of shape [N, 2]")
+                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 must be a torch.Tensor of shape [N, 2]")
+                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 must be a torch.Tensor of shape [N, 2]")
+                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 must be a torch.Tensor of shape [N, 2]")
+                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 must be a torch.Tensor")
+                if pairs.ndim != 2 or pairs.shape[1] != 2:
+                    raise ValueError("pairs must be a torch.Tensor of shape [N, 2]")
+                if pairs.shape[0] == 0:
+                    raise ValueError("pairs length must be greater than 0")
+                if source not in ("timesteps", "sigmas"):
+                    raise ValueError("source must be 'timesteps' or '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 must be positive")
+                    if denoising_strength <= 0:
+                        raise ValueError("denoising_strength must be positive")
+
+                    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 enabled but exponential_shift_mu is missing"
+                            )
+                        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)
+                        if scale_factor != 0:
+                            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:
+        """Return sigma for a scalar timestep via nearest neighbor lookup.
+
+        Args:
+            timestep: Scalar timestep value.
+
+        Returns:
+            Sigma corresponding to the nearest timestep.
+        """
+        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:
+        """Advance one step using an explicit (from, to) timestep pair.
+
+        The update rule is:
+            x_{to} = x_{from} + model_output * (sigma(to) - sigma(from))
+
+        Args:
+            model_output: Predicted model output.
+            timestep_from: Source timestep.
+            timestep_to: Target timestep or None for terminal.
+            sample: Current sample at timestep_from.
+
+        Returns:
+            Updated sample at timestep_to.
+        """
+        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("Scheduler not initialized; call set_timesteps() first")
+
+        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 must return a torch.Tensor")
+            if modified.shape != pairs.shape:
+                raise ValueError("pair_postprocess must return the same shape as input")
+            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/schedulers/hunyuan3d_scheduler.py

@@ -0,0 +1,371 @@
+# Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class Hunyuan3DFlowMatchSchedulerOutput(BaseOutput):
+    """Output class for the scheduler's step function."""
+
+    prev_sample: torch.FloatTensor
+
+
+class Hunyuan3DFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """Euler discrete scheduler for flow matching."""
+
+    # External module path aliases for compatibility with Hunyuan3D configs
+    _aliases = [
+        "hy3dgen.shapegen.schedulers.FlowMatchEulerDiscreteScheduler",
+        "hy3dshape.schedulers.FlowMatchEulerDiscreteScheduler",
+    ]
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        use_dynamic_shifting: bool = False,
+    ):
+        timesteps = np.linspace(
+            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+        ).copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+        sigmas = timesteps / num_train_timesteps
+        if not use_dynamic_shifting:
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
+        self._step_index = None
+        self._begin_index = None
+
+        self.sigmas = sigmas.to("cpu")
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self) -> Optional[int]:
+        """The index counter for current timestep."""
+        return self._step_index
+
+    @property
+    def begin_index(self) -> Optional[int]:
+        """The index for the first timestep."""
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        """Set the begin index for the scheduler.
+
+        Args:
+            begin_index: The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def scale_model_input(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Optional[Union[float, torch.FloatTensor]] = None,
+    ) -> torch.FloatTensor:
+        """Identity operation for flow matching (no input scaling needed)."""
+        return sample
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        noise: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """Forward process in flow-matching (add noise to sample)."""
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timestep
+            ]
+        elif self.step_index is not None:
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+        return sample
+
+    def _sigma_to_t(self, sigma: float) -> float:
+        """Convert sigma to timestep."""
+        return sigma * self.config.num_train_timesteps
+
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        """Apply time shift transformation."""
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[float] = None,
+    ):
+        """Set the discrete timesteps for the diffusion chain."""
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                "Must pass a value for `mu` when `use_dynamic_shifting` is True"
+            )
+
+        if sigmas is None:
+            self.num_inference_steps = num_inference_steps
+            timesteps = np.linspace(
+                self._sigma_to_t(self.sigma_max),
+                self._sigma_to_t(self.sigma_min),
+                num_inference_steps,
+            )
+            sigmas = timesteps / self.config.num_train_timesteps
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
+
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+        timesteps = sigmas * self.config.num_train_timesteps
+
+        self.timesteps = timesteps.to(device=device)
+        self.sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
+
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(
+        self, timestep: float, schedule_timesteps: Optional[torch.Tensor] = None
+    ) -> int:
+        """Find the index for a given timestep."""
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+        pos = 1 if len(indices) > 1 else 0
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep: Union[float, torch.Tensor]):
+        """Initialize step index from timestep."""
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[Hunyuan3DFlowMatchSchedulerOutput, Tuple]:
+        """Predict the sample from the previous timestep."""
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError(
+                "Passing integer indices as timesteps is not supported. "
+                "Pass one of `scheduler.timesteps` as a timestep."
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues
+        sample = sample.to(torch.float32)
+
+        sigma = self.sigmas[self.step_index]
+        sigma_next = self.sigmas[self.step_index + 1]
+
+        prev_sample = sample + (sigma_next - sigma) * model_output
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return Hunyuan3DFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self) -> int:
+        return self.config.num_train_timesteps
+
+
+@dataclass
+class Hunyuan3DConsistencyFlowMatchSchedulerOutput(BaseOutput):
+    """Output for consistency flow matching scheduler."""
+
+    prev_sample: torch.FloatTensor
+    pred_original_sample: torch.FloatTensor
+
+
+class Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """Consistency Flow Matching Euler Discrete Scheduler."""
+
+    # External module path aliases for compatibility with Hunyuan3D configs
+    _aliases = [
+        "hy3dshape.schedulers.Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler",
+    ]
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        pcm_timesteps: int = 50,
+    ):
+        sigmas = np.linspace(0, 1, num_train_timesteps)
+        step_ratio = num_train_timesteps // pcm_timesteps
+
+        euler_timesteps = (np.arange(1, pcm_timesteps) * step_ratio).round().astype(
+            np.int64
+        ) - 1
+        euler_timesteps = np.asarray([0] + euler_timesteps.tolist())
+
+        self.euler_timesteps = euler_timesteps
+        self.sigmas = sigmas[self.euler_timesteps]
+        self.sigmas = torch.from_numpy(self.sigmas.copy()).to(dtype=torch.float32)
+        self.timesteps = self.sigmas * num_train_timesteps
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to("cpu")
+
+    @property
+    def step_index(self) -> Optional[int]:
+        return self._step_index
+
+    @property
+    def begin_index(self) -> Optional[int]:
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        self._begin_index = begin_index
+
+    def scale_model_input(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Optional[Union[float, torch.FloatTensor]] = None,
+    ) -> torch.FloatTensor:
+        """Identity operation for flow matching (no input scaling needed)."""
+        return sample
+
+    def _sigma_to_t(self, sigma: float) -> float:
+        return sigma * self.config.num_train_timesteps
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+    ):
+        """Set timesteps for inference."""
+        self.num_inference_steps = (
+            num_inference_steps if num_inference_steps is not None else len(sigmas)
+        )
+        inference_indices = np.linspace(
+            0, self.config.pcm_timesteps, num=self.num_inference_steps, endpoint=False
+        )
+        inference_indices = np.floor(inference_indices).astype(np.int64)
+        inference_indices = torch.from_numpy(inference_indices).long()
+
+        self.sigmas_ = self.sigmas[inference_indices]
+        timesteps = self.sigmas_ * self.config.num_train_timesteps
+        self.timesteps = timesteps.to(device=device)
+        self.sigmas_ = torch.cat(
+            [self.sigmas_, torch.ones(1, device=self.sigmas_.device)]
+        )
+
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(
+        self, timestep: float, schedule_timesteps: Optional[torch.Tensor] = None
+    ) -> int:
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+        indices = (schedule_timesteps == timestep).nonzero()
+        pos = 1 if len(indices) > 1 else 0
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep: Union[float, torch.Tensor]):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[Hunyuan3DConsistencyFlowMatchSchedulerOutput, Tuple]:
+        """Perform one step of the consistency flow matching scheduler."""
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError("Passing integer indices as timesteps is not supported.")
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        sample = sample.to(torch.float32)
+
+        sigma = self.sigmas_[self.step_index]
+        sigma_next = self.sigmas_[self.step_index + 1]
+
+        prev_sample = sample + (sigma_next - sigma) * model_output
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        pred_original_sample = sample + (1.0 - sigma) * model_output
+        pred_original_sample = pred_original_sample.to(model_output.dtype)
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return Hunyuan3DConsistencyFlowMatchSchedulerOutput(
+            prev_sample=prev_sample, pred_original_sample=pred_original_sample
+        )
+
+    def __len__(self) -> int:
+        return self.config.num_train_timesteps
+
+
+# Entry class for model registry
+EntryClass = [
+    Hunyuan3DFlowMatchEulerDiscreteScheduler,
+    Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler,
+]

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

@@ -0,0 +1,193 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+"""
+Pass-through scheduler for ComfyUI integration.
+
+This scheduler does not modify latents - it simply returns the input sample unchanged.
+The actual denoising logic is handled by ComfyUI.
+"""
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class ComfyUIPassThroughSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor`): The input sample unchanged (pass-through).
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class ComfyUIPassThroughScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
+    """
+    Pass-through scheduler for ComfyUI integration.
+
+    This scheduler does not modify latents. It is used when the denoising logic
+    is handled externally by ComfyUI. The scheduler simply returns the input
+    sample unchanged, allowing ComfyUI to manage the denoising process.
+
+    Usage:
+        - num_inference_steps is always 1 (each step is handled separately)
+        - timesteps are provided externally by ComfyUI
+        - step() returns the input sample unchanged
+    """
+
+    config_name = "scheduler_config.json"
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps=1000,
+        *args,
+        **kwargs,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        # Initialize timesteps as empty - will be set externally
+        self.timesteps = torch.tensor([], dtype=torch.long)
+        self.shift = 0.0
+        self._step_index = 0  # Track current step index
+        self._begin_index: int | None = None  # For compatibility with DenoisingStage
+
+    def set_timesteps(
+        self,
+        num_inference_steps=1,  # Always 1 for ComfyUI
+        timesteps=None,  # Can be provided externally
+        device=None,
+        **kwargs,
+    ):
+        """
+        Set timesteps. For ComfyUI, timesteps are provided externally.
+
+        Args:
+            num_inference_steps: Ignored (always 1 for ComfyUI)
+            timesteps: External timesteps provided by ComfyUI
+            device: Device to place timesteps on
+        """
+        if timesteps is not None:
+            # Use externally provided timesteps
+            if isinstance(timesteps, torch.Tensor):
+                self.timesteps = timesteps
+            else:
+                self.timesteps = torch.tensor(timesteps, dtype=torch.long)
+            if device is not None:
+                self.timesteps = self.timesteps.to(device)
+        else:
+            # Create a single timestep if none provided
+            if device is None:
+                device = torch.device("cpu")
+            self.timesteps = torch.tensor([0], dtype=torch.long, device=device)
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: torch.FloatTensor | int,
+        sample: torch.FloatTensor,
+        return_dict: bool = False,
+        **kwargs,
+    ) -> tuple | ComfyUIPassThroughSchedulerOutput:
+        """
+        Pass-through step: returns the input sample unchanged.
+
+        This scheduler does not modify latents. The actual denoising is handled
+        by ComfyUI, so we simply return the input sample as-is.
+
+        Args:
+            model_output: Predicted noise (ignored, but kept for API compatibility)
+            timestep: Current timestep (ignored, but kept for API compatibility)
+            sample: Input latents (returned unchanged)
+            return_dict: Whether to return a dict or tuple
+
+        Returns:
+            The input sample unchanged (prev_sample = sample)
+        """
+        # Increment step index for tracking
+        self._step_index += 1
+
+        # Simply return the input sample unchanged
+        prev_sample = sample
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return ComfyUIPassThroughSchedulerOutput(prev_sample=prev_sample)
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        """
+        Scale model input. For pass-through scheduler, returns input unchanged.
+
+        Args:
+            sample: Input sample
+            timestep: Timestep (ignored)
+
+        Returns:
+            Input sample unchanged
+        """
+        return sample
+
+    def set_shift(self, shift: float) -> None:
+        """
+        Set shift parameter (no-op for pass-through scheduler).
+
+        Args:
+            shift: Shift value (ignored)
+        """
+        self.shift = shift
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index: The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    @property
+    def begin_index(self) -> int | None:
+        """
+        The index for the first timestep.
+        """
+        return self._begin_index
+
+    @property
+    def step_index(self) -> int:
+        """
+        The index counter for current timestep.
+        """
+        return self._step_index
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Add noise to samples. For pass-through scheduler, returns original samples.
+
+        Args:
+            original_samples: Original clean samples
+            noise: Noise to add (ignored)
+            timestep: Timestep (ignored)
+
+        Returns:
+            Original samples unchanged
+        """
+        return original_samples
+
+
+EntryClass = ComfyUIPassThroughScheduler

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

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# DPM-Solver++ multistep scheduler wrapper for SANA.
+#
+# SANA uses DPM-Solver++ (Lu et al., 2022) as its noise scheduler, which
+# is a high-order ODE solver that converges in fewer steps than DDIM.
+# With solver_order=2 and 20 steps, SANA achieves high-quality results.
+#
+# This wrapper delegates all numerical work to diffusers' implementation
+# and only adapts the interface for sglang's denoising stage.
+
+import torch
+from diffusers import (
+    DPMSolverMultistepScheduler as DiffusersDPMSolverMultistepScheduler,
+)
+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 DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
+    """DPM-Solver++ multistep scheduler wrapper for sglang's BaseScheduler interface."""
+
+    order = 1
+    num_train_timesteps = 1000
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "scaled_linear",
+        trained_betas=None,
+        solver_order: int = 2,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        lower_order_final: bool = True,
+        euler_at_final: bool = False,
+        use_karras_sigmas: bool = False,
+        use_lu_lambdas: bool = False,
+        use_exponential_sigmas: bool = False,
+        use_beta_sigmas: bool = False,
+        use_flow_sigmas: bool = False,
+        final_sigmas_type: str = "zero",
+        lambda_min_clipped: float = -float("inf"),
+        variance_type: str | None = None,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        rescale_betas_zero_snr: bool = False,
+        flow_shift: float | None = None,
+        **kwargs,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self._inner = DiffusersDPMSolverMultistepScheduler(
+            num_train_timesteps=num_train_timesteps,
+            beta_start=beta_start,
+            beta_end=beta_end,
+            beta_schedule=beta_schedule,
+            trained_betas=trained_betas,
+            solver_order=solver_order,
+            prediction_type=prediction_type,
+            thresholding=thresholding,
+            dynamic_thresholding_ratio=dynamic_thresholding_ratio,
+            sample_max_value=sample_max_value,
+            algorithm_type=algorithm_type,
+            solver_type=solver_type,
+            lower_order_final=lower_order_final,
+            euler_at_final=euler_at_final,
+            use_karras_sigmas=use_karras_sigmas,
+            use_lu_lambdas=use_lu_lambdas,
+            use_exponential_sigmas=use_exponential_sigmas,
+            use_beta_sigmas=use_beta_sigmas,
+            use_flow_sigmas=use_flow_sigmas,
+            flow_shift=flow_shift,
+            final_sigmas_type=final_sigmas_type,
+            lambda_min_clipped=lambda_min_clipped,
+            variance_type=variance_type,
+            timestep_spacing=timestep_spacing,
+            steps_offset=steps_offset,
+            rescale_betas_zero_snr=rescale_betas_zero_snr,
+        )
+        self.timesteps = self._inner.timesteps
+        self.order = solver_order
+        self._flow_shift = flow_shift
+        self._begin_index: int | None = None
+        BaseScheduler.__init__(self)
+
+    def set_shift(self, shift: float) -> None:
+        self._flow_shift = shift
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        self._begin_index = begin_index
+
+    @property
+    def begin_index(self) -> int | None:
+        return self._begin_index
+
+    def set_timesteps(self, num_inference_steps: int, device=None, **kwargs):
+        self._inner.set_timesteps(num_inference_steps, device=device, **kwargs)
+        self.timesteps = self._inner.timesteps
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        return self._inner.scale_model_input(sample, timestep)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int,
+        sample: torch.Tensor,
+        generator: torch.Generator | None = None,
+        variance_noise: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ):
+        return self._inner.step(
+            model_output,
+            timestep,
+            sample,
+            generator=generator,
+            variance_noise=variance_noise,
+            return_dict=return_dict,
+        )
+
+    @property
+    def sigmas(self):
+        return getattr(self._inner, "sigmas", None)
+
+    @property
+    def init_noise_sigma(self):
+        return self._inner.init_noise_sigma
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        return self._inner.add_noise(original_samples, noise, timesteps)
+
+
+EntryClass = DPMSolverMultistepScheduler

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

@@ -0,0 +1,703 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import scipy.stats
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler
+from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import (
+    SchedulerRLMixin,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchEulerDiscreteScheduler(
+    SchedulerMixin, ConfigMixin, BaseScheduler, SchedulerRLMixin
+):
+    """
+    Euler scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+        use_dynamic_shifting (`bool`, defaults to False):
+            Whether to apply timestep shifting on-the-fly based on the image resolution.
+        base_shift (`float`, defaults to 0.5):
+            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
+            with desired output.
+        max_shift (`float`, defaults to 1.15):
+            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
+            more exaggerated or stylized.
+        base_image_seq_len (`int`, defaults to 256):
+            The base image sequence length.
+        max_image_seq_len (`int`, defaults to 4096):
+            The maximum image sequence length.
+        invert_sigmas (`bool`, defaults to False):
+            Whether to invert the sigmas.
+        shift_terminal (`float`, defaults to None):
+            The end value of the shifted timestep schedule.
+        use_karras_sigmas (`bool`, defaults to False):
+            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
+        use_exponential_sigmas (`bool`, defaults to False):
+            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
+        use_beta_sigmas (`bool`, defaults to False):
+            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
+        time_shift_type (`str`, defaults to "exponential"):
+            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
+        stochastic_sampling (`bool`, defaults to False):
+            Whether to use stochastic sampling.
+    """
+
+    _compatibles: list[Any] = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        use_dynamic_shifting: bool = False,
+        base_shift: float | None = 0.5,
+        max_shift: float | None = 1.15,
+        base_image_seq_len: int | None = 256,
+        max_image_seq_len: int | None = 4096,
+        invert_sigmas: bool = False,
+        shift_terminal: float | None = None,
+        use_karras_sigmas: bool | None = False,
+        use_exponential_sigmas: bool | None = False,
+        use_beta_sigmas: bool | None = False,
+        time_shift_type: str = "exponential",
+        stochastic_sampling: bool = False,
+    ):
+        if (
+            sum(
+                [
+                    self.config.use_beta_sigmas,
+                    self.config.use_exponential_sigmas,
+                    self.config.use_karras_sigmas,
+                ]
+            )
+            > 1
+        ):
+            raise ValueError(
+                "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
+            )
+        if time_shift_type not in {"exponential", "linear"}:
+            raise ValueError(
+                "`time_shift_type` must either be 'exponential' or 'linear'."
+            )
+
+        timesteps = np.linspace(
+            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+        )[::-1].copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+        sigmas = timesteps / num_train_timesteps
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
+        self.num_train_timesteps = num_train_timesteps
+
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
+
+        self._shift = shift
+
+        self.sigmas = sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+        BaseScheduler.__init__(self)
+
+    @property
+    def shift(self) -> float:
+        """
+        The value used for shifting.
+        """
+        return self._shift
+
+    @property
+    def step_index(self) -> int | None:
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self) -> int | None:
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def set_shift(self, shift: float) -> None:
+        self._shift = shift
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        noise: torch.FloatTensor | None = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward process in flow-matching
+        """
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+            assert isinstance(timestep, torch.Tensor)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            assert isinstance(timestep, torch.Tensor)
+            timestep = timestep.to(sample.device)
+
+        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timestep
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+
+        return sample
+
+    def _sigma_to_t(self, sigma: float) -> float:
+        return sigma * self.config.num_train_timesteps
+
+    def time_shift(
+        self, mu: float, sigma: float, t: torch.Tensor | np.ndarray
+    ) -> torch.Tensor | np.ndarray:
+        if self.config.time_shift_type == "exponential":
+            return self._time_shift_exponential(mu, sigma, t)
+        elif self.config.time_shift_type == "linear":
+            return self._time_shift_linear(mu, sigma, t)
+        else:
+            raise ValueError(f"Unknown time_shift_type: {self.config.time_shift_type}")
+
+    def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
+        r"""
+        Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
+        value.
+
+        Reference:
+        https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51
+
+        Args:
+            t (`torch.Tensor`):
+                A tensor of timesteps to be stretched and shifted.
+
+        Returns:
+            `torch.Tensor`:
+                A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
+        """
+        one_minus_z = 1 - t
+        scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
+        if scale_factor == 0:
+            return t
+        stretched_t = 1 - (one_minus_z / scale_factor)
+        return stretched_t
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device = None,
+        sigmas: list[float] | None = None,
+        mu: float | None = None,
+        timesteps: list[float] | None = None,
+    ) -> None:
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`, *optional*):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            sigmas (`List[float]`, *optional*):
+                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
+                automatically.
+            mu (`float`, *optional*):
+                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
+                shifting.
+            timesteps (`List[float]`, *optional*):
+                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
+                automatically.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                "`mu` must be passed when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if (
+            sigmas is not None
+            and timesteps is not None
+            and len(sigmas) != len(timesteps)
+        ):
+            raise ValueError("`sigmas` and `timesteps` should have the same length")
+
+        if num_inference_steps is not None:
+            if (sigmas is not None and len(sigmas) != num_inference_steps) or (
+                timesteps is not None and len(timesteps) != num_inference_steps
+            ):
+                raise ValueError(
+                    "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided"
+                )
+        else:
+            if sigmas is not None:
+                num_inference_steps = len(sigmas)
+            elif timesteps is not None:
+                num_inference_steps = len(timesteps)
+            else:
+                raise ValueError(
+                    "Either num_inference_steps, sigmas, or timesteps must be provided"
+                )
+
+        self.num_inference_steps = num_inference_steps
+
+        # 1. Prepare default sigmas
+        is_timesteps_provided = timesteps is not None
+
+        timesteps_array: np.ndarray | None = None
+        if is_timesteps_provided:
+            assert timesteps is not None
+            timesteps_array = np.array(timesteps).astype(np.float32)
+
+        sigmas_array: np.ndarray
+        if sigmas is None:
+            if timesteps_array is None:
+                timesteps_array = np.linspace(
+                    self._sigma_to_t(self.sigma_max),
+                    self._sigma_to_t(self.sigma_min),
+                    num_inference_steps,
+                )
+            sigmas_array = timesteps_array / self.config.num_train_timesteps
+        else:
+            sigmas_array = np.array(sigmas).astype(np.float32)
+            num_inference_steps = len(sigmas_array)
+
+        # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
+        #    "exponential" or "linear" type is applied
+        if self.config.use_dynamic_shifting:
+            assert mu is not None, "mu cannot be None when use_dynamic_shifting is True"
+            sigmas_array = self.time_shift(mu, 1.0, sigmas_array)
+        else:
+            sigmas_array = (
+                self.shift * sigmas_array / (1 + (self.shift - 1) * sigmas_array)
+            )
+
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
+        if self.config.shift_terminal:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32)
+            sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor)
+            sigmas_array = sigmas_tensor.numpy()
+
+        # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
+        if self.config.use_karras_sigmas:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32)
+            sigmas_tensor = self._convert_to_karras(
+                in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps
+            )
+            sigmas_array = sigmas_tensor.numpy()
+        elif self.config.use_exponential_sigmas:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32)
+            sigmas_tensor = self._convert_to_exponential(
+                in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps
+            )
+            sigmas_array = sigmas_tensor.numpy()
+        elif self.config.use_beta_sigmas:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32)
+            sigmas_tensor = self._convert_to_beta(
+                in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps
+            )
+            sigmas_array = sigmas_tensor.numpy()
+
+        # 5. Convert sigmas and timesteps to tensors and move to specified device
+        sigmas_tensor = torch.from_numpy(sigmas_array).to(
+            dtype=torch.float32, device=device
+        )
+        if not is_timesteps_provided:
+            timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps
+        else:
+            assert timesteps_array is not None
+            timesteps_tensor = torch.from_numpy(timesteps_array).to(
+                dtype=torch.float32, device=device
+            )
+
+        # 6. Append the terminal sigma value.
+        #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
+        #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
+        if self.config.invert_sigmas:
+            sigmas_tensor = 1.0 - sigmas_tensor
+            timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps
+            sigmas_tensor = torch.cat(
+                [sigmas_tensor, torch.ones(1, device=sigmas_tensor.device)]
+            )
+        else:
+            sigmas_tensor = torch.cat(
+                [sigmas_tensor, torch.zeros(1, device=sigmas_tensor.device)]
+            )
+
+        self.timesteps = timesteps_tensor
+        self.sigmas = sigmas_tensor
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(
+        self,
+        timestep: float | torch.FloatTensor,
+        schedule_timesteps: torch.Tensor | None = None,
+    ) -> int:
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep: float | torch.FloatTensor) -> None:
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int | torch.Tensor,
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        batch=None,
+        return_dict: bool = True,
+    ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`int` or `torch.Tensor`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
+            return_dict (`bool`):
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
+
+        Returns:
+            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`,
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
+                otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        if isinstance(timestep, int | torch.IntTensor | torch.LongTensor):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+
+        if per_token_timesteps is not None:
+            per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps
+
+            sigmas = self.sigmas[:, None, None]
+            lower_mask = sigmas < per_token_sigmas[None] - 1e-6
+            lower_sigmas = lower_mask * sigmas
+            lower_sigmas, _ = lower_sigmas.max(dim=0)
+
+            current_sigma = per_token_sigmas[..., None]
+            next_sigma = lower_sigmas[..., None]
+            dt = current_sigma - next_sigma
+        else:
+            assert self.step_index is not None, "step_index should not be None"
+            sigma_idx = self.step_index
+            sigma = self.sigmas[sigma_idx]
+            sigma_next = self.sigmas[sigma_idx + 1]
+
+            current_sigma = sigma
+            next_sigma = sigma_next
+            dt = sigma_next - sigma
+
+        if batch is not None and batch.rollout:
+            if not self.already_prepared_rollout(batch):
+                raise RuntimeError("Rollout not prepared before step")
+            prev_sample = self.flow_sde_sampling(
+                batch, model_output, sample, current_sigma, next_sigma, generator
+            )
+        else:
+            if self.config.stochastic_sampling:
+                x0 = sample - current_sigma * model_output
+                noise = torch.randn_like(sample)
+                prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
+            else:
+                prev_sample = sample + dt * model_output
+
+        # upon completion increase step index by one
+        assert self._step_index is not None, "_step_index should not be None"
+        self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
+
+        if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
+            return (prev_sample,)
+
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(
+        self, in_sigmas: torch.Tensor, num_inference_steps: int
+    ) -> torch.Tensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
+    def _convert_to_exponential(
+        self, in_sigmas: torch.Tensor, num_inference_steps: int
+    ) -> torch.Tensor:
+        """Constructs an exponential noise schedule."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.exp(
+            np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)
+        )
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
+    def _convert_to_beta(
+        self,
+        in_sigmas: torch.Tensor,
+        num_inference_steps: int,
+        alpha: float = 0.6,
+        beta: float = 0.6,
+    ) -> torch.Tensor:
+        """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.array(
+            [
+                sigma_min + (ppf * (sigma_max - sigma_min))
+                for ppf in [
+                    scipy.stats.beta.ppf(timestep, alpha, beta)
+                    for timestep in 1 - np.linspace(0, 1, num_inference_steps)
+                ]
+            ]
+        )
+        return sigmas
+
+    def _time_shift_exponential(
+        self, mu: float, sigma: float, t: torch.Tensor | np.ndarray
+    ) -> torch.Tensor | np.ndarray:
+        if isinstance(t, np.ndarray):
+            return np.exp(mu) / (np.exp(mu) + (1 / t - 1) ** sigma)
+        else:
+            return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def _time_shift_linear(
+        self, mu: float, sigma: float, t: torch.Tensor | np.ndarray
+    ) -> torch.Tensor | np.ndarray:
+        return mu / (mu + (1 / t - 1) ** sigma)
+
+    def add_noise(
+        self,
+        clean_latent: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.IntTensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            clean_latent: the clean latent with shape [B, C, H, W],
+                where B is batch_size or batch_size * num_frames
+            noise: the noise with shape [B, C, H, W]
+            timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
+
+        Returns:
+            the corrupted latent with shape [B, C, H, W]
+        """
+        # If timestep is [bs, num_frames]
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+            assert timestep.numel() == clean_latent.shape[0]
+        elif timestep.ndim == 1:
+            # If timestep is [1]
+            if timestep.shape[0] == 1:
+                timestep = timestep.expand(clean_latent.shape[0])
+            else:
+                assert timestep.numel() == clean_latent.shape[0]
+        else:
+            raise ValueError(f"[add_noise] Invalid timestep shape: {timestep.shape}")
+        # timestep shape should be [B]
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
+        )
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * clean_latent + sigma * noise
+        return sample.type_as(noise)
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        return sample
+
+    def __len__(self) -> int:
+        return 0
+
+
+EntryClass = FlowMatchEulerDiscreteScheduler

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

@@ -0,0 +1,843 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
+# Convert unipc for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import math
+from typing import Any
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (
+    KarrasDiffusionSchedulers,
+    SchedulerMixin,
+    SchedulerOutput,
+)
+from diffusers.utils import deprecate
+
+from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler
+
+
+class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
+    """
+    `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        solver_order (`int`, default `2`):
+            The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
+            due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
+            unconditional sampling.
+        prediction_type (`str`, defaults to "flow_prediction"):
+            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+            the flow of the diffusion process.
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
+        predict_x0 (`bool`, defaults to `True`):
+            Whether to use the updating algorithm on the predicted x0.
+        solver_type (`str`, default `bh2`):
+            Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
+            otherwise.
+        lower_order_final (`bool`, default `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        disable_corrector (`list`, default `[]`):
+            Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
+            and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
+            usually disabled during the first few steps.
+        solver_p (`SchedulerMixin`, default `None`):
+            Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        final_sigmas_type (`str`, defaults to `"zero"`):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        solver_order: int = 2,
+        prediction_type: str = "flow_prediction",
+        shift: float | None = 1.0,
+        use_dynamic_shifting=False,
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        predict_x0: bool = True,
+        solver_type: str = "bh2",
+        lower_order_final: bool = True,
+        disable_corrector: tuple = (),
+        solver_p: SchedulerMixin = None,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        final_sigmas_type: str | None = "zero",  # "zero", "sigma_min"
+        **kwargs,
+    ):
+
+        if solver_type not in ["bh1", "bh2"]:
+            if solver_type in ["midpoint", "heun", "logrho"]:
+                self.register_to_config(solver_type="bh2")
+            else:
+                raise NotImplementedError(
+                    f"{solver_type} is not implemented for {self.__class__}"
+                )
+
+        self.predict_x0 = predict_x0
+        # setable values
+        self.num_inference_steps: int | None = None
+        alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[
+            ::-1
+        ].copy()
+        sigmas = 1.0 - alphas
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            assert shift is not None
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
+
+        self.sigmas = sigmas
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+        self.timesteps = sigmas * num_train_timesteps
+        self.num_train_timesteps = num_train_timesteps
+
+        self.model_outputs = [None] * solver_order
+        self.timestep_list: list[Any | None] = [None] * solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = list(disable_corrector)
+        self.solver_p = solver_p
+        self.last_sample = None
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
+
+        BaseScheduler.__init__(self)
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    def set_shift(self, shift: float) -> None:
+        self.config.shift = shift
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device = None,
+        sigmas: list[float] | None = None,
+        mu: float | None | None = None,
+        shift: float | None | None = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                Total number of the spacing of the time steps.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            assert num_inference_steps is not None
+            sigmas = np.linspace(
+                self.sigma_max, self.sigma_min, num_inference_steps + 1
+            ).copy()[
+                :-1
+            ]  # pyright: ignore
+
+        if self.config.use_dynamic_shifting:
+            assert mu is not None
+            sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
+        else:
+            if shift is None:
+                shift = self.config.shift
+            assert isinstance(sigmas, np.ndarray)
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = sigmas[-1]
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        timesteps = sigmas * self.config.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [sigma_last]]).astype(
+            np.float32
+        )  # pyright: ignore
+
+        self.sigmas = torch.from_numpy(sigmas).to(device=device)
+        self.timesteps = torch.from_numpy(timesteps).to(
+            device=device, dtype=torch.int64
+        )
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+        self.last_sample = None
+        if self.solver_p:
+            self.solver_p.set_timesteps(self.num_inference_steps, device=device)
+
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+        self._begin_index = None
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = (
+                sample.float()
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = (
+            torch.clamp(sample, -s, s) / s
+        )  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    def _sigma_to_alpha_sigma_t(self, sigma) -> tuple[Any, Any]:
+        return 1 - sigma, sigma
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Convert the model output to the corresponding type the UniPC algorithm needs.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError("missing `sample` as a required keyword argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if self.predict_x0:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+        else:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                epsilon = sample - (1 - sigma_t) * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = model_output + x0_pred
+
+            return epsilon
+
+    def multistep_uni_p_bh_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        order: int | None = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model at the current timestep.
+            prev_timestep (`int`):
+                The previous discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            order (`int`):
+                The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(" missing `sample` as a required keyword argument")
+        if order is None:
+            if len(args) > 2:
+                order = args[2]
+            else:
+                raise ValueError(" missing `order` as a required keyword argument")
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+        model_output_list = self.model_outputs
+
+        s0 = self.timestep_list[-1]
+        m0 = model_output_list[-1]
+        x = sample
+
+        if self.solver_p:
+            x_t = self.solver_p.step(model_output, s0, x).prev_sample
+            return x_t
+
+        sigma_t, sigma_s0 = (
+            self.sigmas[self.step_index + 1],
+            self.sigmas[self.step_index],
+        )  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = sample.device
+
+        rks = []
+        D1s: list[Any] | None = []
+        sigmas = self.sigmas.to(device=device)
+        for i in range(1, order):
+            si = self.step_index - i  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            assert mi is not None
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        if len(rks) > 0:
+            rks = torch.stack(rks)
+            one = torch.ones(1, device=device, dtype=rks.dtype)
+            rks = torch.cat([rks, one])
+        else:
+            rks = torch.ones(1, device=device, dtype=h.dtype)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.stack(b)
+
+        if D1s is not None and len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            # for order 2, we use a simplified version
+            if order == 2:
+                rhos_p = 0.5 * torch.ones(1, dtype=x.dtype, device=device)
+            else:
+                assert isinstance(R, torch.Tensor)
+                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
+        else:
+            D1s = None
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum(
+                    "k,bkc...->bc...", rhos_p, D1s
+                )  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - alpha_t * B_h * pred_res
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum(
+                    "k,bkc...->bc...", rhos_p, D1s
+                )  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - sigma_t * B_h * pred_res
+
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output: torch.Tensor,
+        *args,
+        last_sample: torch.Tensor = None,
+        this_sample: torch.Tensor = None,
+        order: int | None = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniC (B(h) version).
+
+        Args:
+            this_model_output (`torch.Tensor`):
+                The model outputs at `x_t`.
+            this_timestep (`int`):
+                The current timestep `t`.
+            last_sample (`torch.Tensor`):
+                The generated sample before the last predictor `x_{t-1}`.
+            this_sample (`torch.Tensor`):
+                The generated sample after the last predictor `x_{t}`.
+            order (`int`):
+                The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+
+        Returns:
+            `torch.Tensor`:
+                The corrected sample tensor at the current timestep.
+        """
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None)
+        if last_sample is None:
+            if len(args) > 1:
+                last_sample = args[1]
+            else:
+                raise ValueError(" missing`last_sample` as a required keyword argument")
+        if this_sample is None:
+            if len(args) > 2:
+                this_sample = args[2]
+            else:
+                raise ValueError(" missing`this_sample` as a required keyword argument")
+        if order is None:
+            if len(args) > 3:
+                order = args[3]
+            else:
+                raise ValueError(" missing`order` as a required keyword argument")
+        if this_timestep is not None:
+            deprecate(
+                "this_timestep",
+                "1.0.0",
+                "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        model_output_list = self.model_outputs
+
+        m0 = model_output_list[-1]
+        x = last_sample
+        x_t = this_sample
+        model_t = this_model_output
+
+        sigma_t, sigma_s0 = (
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],
+        )  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+
+        # Build rks and D1s fully on device to avoid any host-device sync
+        # Fast paths for small orders (common cases: 1 or 2)
+        if order == 1:
+            rks = torch.ones(1, device=device, dtype=h.dtype)
+            D1s = None
+        elif order == 2:
+            # order == 2 -> only one historical point is used
+            si = self.step_index - 2  # i = 1
+            mi = model_output_list[-2]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h  # 0-dim tensor on device
+            # rks = [rk, 1.0] but keep it on device without list->tensor sync
+            rks = torch.stack((rk, torch.ones_like(rk)))
+            assert mi is not None
+            # D1s shape: (B, K=1, C, ...) to match later einsum over K
+            D1s = ((mi - m0) / rk).unsqueeze(1)  # pyright: ignore
+        else:
+            rks_list = []
+            D1s_list = []
+            for i in range(1, order):
+                si = self.step_index - (i + 1)
+                mi = model_output_list[-(i + 1)]
+                alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+                lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+                rk = (lambda_si - lambda_s0) / h
+                rks_list.append(rk)
+                assert mi is not None
+                D1s_list.append((mi - m0) / rk)  # pyright: ignore
+
+            # Append 1.0 as a device tensor to rks
+            rks = torch.stack(rks_list + [torch.ones_like(rks_list[0])])
+            D1s = torch.stack(D1s_list, dim=1) if len(D1s_list) > 0 else None
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        # Avoid torch.tensor(list_of_gpu_scalars) which syncs to host
+        b = torch.stack(b)
+
+        # D1s is already prepared above for order==2; remains None for order==1
+
+        # for order 1, we use a simplified version
+        if order == 1:
+            rhos_c = 0.5 * torch.ones(1, dtype=x.dtype, device=device)
+        elif order == 2:
+            # Manually solve the 2x2 linear system to avoid device synchronization from torch.linalg.solve
+            # R = [[1, 1], [rk, 1]], where rk = rks[0]
+            rk = rks[0]
+            det = 1 - rk
+            # Using Cramer's rule to solve for rhos_c = [x0, x1]
+            # x0 = (b0 - b1) / det
+            # x1 = (b1 - rk * b0) / det
+            rhos_c_0 = (b[0] - b[1]) / det
+            rhos_c_1 = (b[1] - rk * b[0]) / det
+            rhos_c = torch.stack([rhos_c_0, rhos_c_1])
+        else:
+            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+        step_index: int = indices[pos].item()
+
+        return step_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
+    def _init_step_index(self, timestep) -> None:
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int | torch.Tensor,
+        sample: torch.Tensor,
+        return_dict: bool = True,
+        generator=None,
+    ) -> SchedulerOutput | tuple:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep UniPC.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to call 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        use_corrector = (
+            self.step_index > 0
+            and self.step_index - 1 not in self.disable_corrector
+            and self.last_sample is not None  # pyright: ignore
+        )
+
+        sample = sample.to(model_output.device)
+        model_output_convert = self.convert_model_output(model_output, sample=sample)
+
+        if use_corrector:
+            sample = self.multistep_uni_c_bh_update(
+                this_model_output=model_output_convert,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self.timestep_list[i] = self.timestep_list[i + 1]
+
+        self.model_outputs[-1] = model_output_convert
+        self.timestep_list[-1] = timestep  # pyright: ignore
+
+        if self.config.lower_order_final:
+            this_order = min(
+                self.config.solver_order, len(self.timesteps) - self.step_index
+            )  # pyright: ignore
+        else:
+            this_order = self.config.solver_order
+
+        self.this_order: int = min(
+            this_order, self.lower_order_nums + 1
+        )  # warmup for multistep
+        assert self.this_order > 0
+
+        self.last_sample = sample
+        prev_sample = self.multistep_uni_p_bh_update(
+            model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
+            sample=sample,
+            order=self.this_order,
+        )
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # upon completion increase step index by one
+        assert self._step_index is not None
+        self._step_index += 1  # pyright: ignore
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(
+            device=original_samples.device, dtype=original_samples.dtype
+        )
+        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(
+                original_samples.device, dtype=torch.float32
+            )
+            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timesteps
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timesteps.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+
+
+EntryClass = FlowUniPCMultistepScheduler

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

@@ -0,0 +1,744 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from Helios diffusers scheduler:
+# https://github.com/BestWishYsh/Helios
+"""
+Helios scheduler implementing flow-matching with UniPC/Euler solvers.
+
+For Phase 1 T2V (stages=1), this simplifies to standard flow-matching
+with dynamic shifting and UniPC multistep solver.
+"""
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+@dataclass
+class HeliosSchedulerOutput:
+    prev_sample: torch.FloatTensor
+    model_outputs: torch.FloatTensor | None = None
+    last_sample: torch.FloatTensor | None = None
+    this_order: int | None = None
+
+
+class HeliosSchedulerConfig:
+    """Mimics diffusers config interface for scheduler parameters."""
+
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
+
+class HeliosScheduler:
+    """
+    Helios multi-stage scheduler supporting Euler, UniPC, and DMD solvers.
+
+    For Phase 1 T2V with stages=1, this is a standard flow-matching scheduler
+    with optional time shifting and UniPC multistep updates.
+    """
+
+    order = 1
+
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        stages: int = 1,
+        stage_range: list | None = None,
+        gamma: float = 1 / 3,
+        thresholding: bool = False,
+        prediction_type: str = "flow_prediction",
+        solver_order: int = 2,
+        predict_x0: bool = True,
+        solver_type: str = "bh2",
+        lower_order_final: bool = True,
+        disable_corrector: list[int] | None = None,
+        use_flow_sigmas: bool = True,
+        scheduler_type: str = "unipc",
+        use_dynamic_shifting: bool = False,
+        time_shift_type: str = "linear",
+        **kwargs,
+    ):
+        if stage_range is None:
+            # Evenly divide [0, 1] into 3 stages for pyramid SR
+            stage_range = [0, 1 / 3, 2 / 3, 1]
+        if disable_corrector is None:
+            disable_corrector = []
+
+        self.config = HeliosSchedulerConfig(
+            num_train_timesteps=num_train_timesteps,
+            shift=shift,
+            stages=stages,
+            stage_range=stage_range,
+            gamma=gamma,
+            thresholding=thresholding,
+            prediction_type=prediction_type,
+            solver_order=solver_order,
+            predict_x0=predict_x0,
+            solver_type=solver_type,
+            lower_order_final=lower_order_final,
+            disable_corrector=disable_corrector,
+            use_flow_sigmas=use_flow_sigmas,
+            scheduler_type=scheduler_type,
+            use_dynamic_shifting=use_dynamic_shifting,
+            time_shift_type=time_shift_type,
+        )
+
+        self.timestep_ratios = {}
+        self.timesteps_per_stage = {}
+        self.sigmas_per_stage = {}
+        self.start_sigmas = {}
+        self.end_sigmas = {}
+        self.ori_start_sigmas = {}
+
+        self.init_sigmas_for_each_stage()
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+        self.gamma = gamma
+
+        if solver_type not in ["bh1", "bh2"]:
+            raise NotImplementedError(f"{solver_type} is not implemented")
+
+        self.predict_x0 = predict_x0
+        self.model_outputs = [None] * solver_order
+        self.timestep_list = [None] * solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = disable_corrector
+        self.solver_p = None
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+        self.num_inference_steps = None
+
+    def init_sigmas(self):
+        num_train_timesteps = self.config.num_train_timesteps
+        shift = self.config.shift
+
+        alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1)
+        sigmas = 1.0 - alphas
+        sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy()
+        sigmas = torch.from_numpy(sigmas)
+        timesteps = (sigmas * num_train_timesteps).clone()
+
+        self._step_index = None
+        self._begin_index = None
+        self.timesteps = timesteps
+        self.sigmas = sigmas.to("cpu")
+
+    def init_sigmas_for_each_stage(self):
+        self.init_sigmas()
+
+        stage_distance = []
+        stages = self.config.stages
+        training_steps = self.config.num_train_timesteps
+        stage_range = self.config.stage_range
+
+        for i_s in range(stages):
+            start_indice = int(stage_range[i_s] * training_steps)
+            start_indice = max(start_indice, 0)
+            end_indice = int(stage_range[i_s + 1] * training_steps)
+            end_indice = min(end_indice, training_steps)
+            start_sigma = self.sigmas[start_indice].item()
+            end_sigma = (
+                self.sigmas[end_indice].item() if end_indice < training_steps else 0.0
+            )
+            self.ori_start_sigmas[i_s] = start_sigma
+
+            if i_s != 0:
+                ori_sigma = 1 - start_sigma
+                gamma = self.config.gamma
+                corrected_sigma = (
+                    1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)
+                ) * ori_sigma
+                start_sigma = 1 - corrected_sigma
+
+            stage_distance.append(start_sigma - end_sigma)
+            self.start_sigmas[i_s] = start_sigma
+            self.end_sigmas[i_s] = end_sigma
+
+        tot_distance = sum(stage_distance)
+        for i_s in range(stages):
+            if i_s == 0:
+                start_ratio = 0.0
+            else:
+                start_ratio = sum(stage_distance[:i_s]) / tot_distance
+            if i_s == stages - 1:
+                # Use value just below 1.0 to avoid out-of-bounds indexing
+                end_ratio = 1.0 - 1e-16
+            else:
+                end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance
+            self.timestep_ratios[i_s] = (start_ratio, end_ratio)
+
+        for i_s in range(stages):
+            timestep_ratio = self.timestep_ratios[i_s]
+            # Clamp to max valid timestep (num_train_timesteps - 1)
+            timestep_max = min(
+                self.timesteps[int(timestep_ratio[0] * training_steps)], 999
+            )
+            timestep_min = self.timesteps[
+                min(int(timestep_ratio[1] * training_steps), training_steps - 1)
+            ]
+            timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1)
+            self.timesteps_per_stage[i_s] = (
+                timesteps[:-1]
+                if isinstance(timesteps, torch.Tensor)
+                else torch.from_numpy(timesteps[:-1])
+            )
+            # Sigma range [0.999, 0]: start just below 1.0 to avoid singularity
+            stage_sigmas = np.linspace(0.999, 0, training_steps + 1)
+            self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1])
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        self._begin_index = begin_index
+
+    def time_shift(self, mu, sigma, t):
+        if self.config.time_shift_type == "exponential":
+            return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+        elif self.config.time_shift_type == "linear":
+            return mu / (mu + (1 / t - 1) ** sigma)
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        stage_index: int | None = None,
+        device: str | torch.device = None,
+        sigmas=None,
+        mu=None,
+        is_amplify_first_chunk: bool = False,
+    ):
+        if self.config.scheduler_type == "dmd":
+            if is_amplify_first_chunk:
+                num_inference_steps = num_inference_steps * 2 + 1
+            else:
+                num_inference_steps = num_inference_steps + 1
+
+        self.num_inference_steps = num_inference_steps
+        self.init_sigmas()
+
+        if self.config.stages == 1:
+            if sigmas is None:
+                sigmas = np.linspace(
+                    1,
+                    1 / self.config.num_train_timesteps,
+                    num_inference_steps + 1,
+                )[:-1].astype(np.float32)
+                if self.config.shift != 1.0:
+                    assert not self.config.use_dynamic_shifting
+                    sigmas = self.time_shift(self.config.shift, 1.0, sigmas)
+            timesteps = (sigmas * self.config.num_train_timesteps).copy()
+            sigmas = torch.from_numpy(sigmas)
+        else:
+            stage_timesteps = self.timesteps_per_stage[stage_index]
+            timesteps = np.linspace(
+                stage_timesteps[0].item(),
+                stage_timesteps[-1].item(),
+                num_inference_steps,
+            )
+            stage_sigmas = self.sigmas_per_stage[stage_index]
+            ratios = np.linspace(
+                stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps
+            )
+            sigmas = torch.from_numpy(ratios)
+
+        self.timesteps = torch.from_numpy(timesteps).to(device=device)
+        self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
+        if current_platform.is_npu():
+            # self.sigmas is float64 (np.linspace default); Ascend aclnnExpm1 does
+            # not support float64 (DT_DOUBLE) and crashes the UniPC step's expm1.
+            # Pin fp32 on NPU; remove once aclnnExpm1 supports float64.
+            self.sigmas = self.sigmas.to(torch.float32)
+
+        self._step_index = None
+        self.reset_scheduler_history()
+
+        if self.config.scheduler_type == "dmd":
+            self.timesteps = self.timesteps[:-1]
+            self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]])
+
+        if self.config.use_dynamic_shifting:
+            assert self.config.shift == 1.0
+            self.sigmas = self.time_shift(mu, 1.0, self.sigmas)
+            if self.config.stages == 1:
+                self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps
+            else:
+                self.timesteps = self.timesteps_per_stage[
+                    stage_index
+                ].min() + self.sigmas[:-1] * (
+                    self.timesteps_per_stage[stage_index].max()
+                    - self.timesteps_per_stage[stage_index].min()
+                )
+
+    # ---------------------------------- Euler ----------------------------------
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+        indices = (schedule_timesteps == timestep).nonzero()
+        pos = 1 if len(indices) > 1 else 0
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step_euler(
+        self,
+        model_output: torch.FloatTensor,
+        timestep=None,
+        sample: torch.FloatTensor = None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        if self.step_index is None:
+            self._step_index = 0
+
+        sample = sample.to(torch.float32)
+        sigma = self.sigmas[self.step_index]
+        sigma_next = self.sigmas[self.step_index + 1]
+
+        prev_sample = sample + (sigma_next - sigma) * model_output
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return HeliosSchedulerOutput(prev_sample=prev_sample)
+
+    # ---------------------------------- UniPC ----------------------------------
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        if self.config.use_flow_sigmas:
+            alpha_t = 1 - sigma
+            sigma_t = torch.clamp(sigma, min=1e-8)
+        else:
+            alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
+            sigma_t = sigma * alpha_t
+        return alpha_t, sigma_t
+
+    def convert_model_output(self, model_output, sample=None, sigma=None, **kwargs):
+        flag = False
+        if sigma is None:
+            flag = True
+            sigma = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+        if self.predict_x0:
+            if self.config.prediction_type == "flow_prediction":
+                if flag:
+                    sigma_t = self.sigmas[self.step_index]
+                else:
+                    sigma_t = sigma
+                x0_pred = sample - sigma_t * model_output
+            elif self.config.prediction_type == "epsilon":
+                x0_pred = (sample - sigma_t * model_output) / alpha_t
+            elif self.config.prediction_type == "sample":
+                x0_pred = model_output
+            elif self.config.prediction_type == "v_prediction":
+                x0_pred = alpha_t * sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type {self.config.prediction_type} not supported"
+                )
+            return x0_pred
+        else:
+            if self.config.prediction_type == "epsilon":
+                return model_output
+            elif self.config.prediction_type == "sample":
+                return (sample - alpha_t * model_output) / sigma_t
+            elif self.config.prediction_type == "v_prediction":
+                return alpha_t * model_output + sigma_t * sample
+            else:
+                raise ValueError(
+                    f"prediction_type {self.config.prediction_type} not supported"
+                )
+
+    def multistep_uni_p_bh_update(
+        self, model_output, sample=None, order=None, sigma=None, sigma_next=None
+    ):
+        model_output_list = self.model_outputs
+        m0 = model_output_list[-1]
+        x = sample
+
+        if sigma_next is None and sigma is None:
+            sigma_t, sigma_s0 = (
+                self.sigmas[self.step_index + 1],
+                self.sigmas[self.step_index],
+            )
+        else:
+            sigma_t, sigma_s0 = sigma_next, sigma
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        h = lambda_t - lambda_s0
+        device = sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - i
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)
+        h_phi_k = h_phi_1 / hh - 1
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+            if order == 2:
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
+        else:
+            D1s = None
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            pred_res = (
+                torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0
+            )
+            x_t = x_t_ - alpha_t * B_h * pred_res
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            pred_res = (
+                torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0
+            )
+            x_t = x_t_ - sigma_t * B_h * pred_res
+
+        return x_t.to(x.dtype)
+
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output,
+        last_sample=None,
+        this_sample=None,
+        order=None,
+        sigma_before=None,
+        sigma=None,
+    ):
+        model_output_list = self.model_outputs
+        m0 = model_output_list[-1]
+        x = last_sample
+        model_t = this_model_output
+
+        if sigma_before is None and sigma is None:
+            sigma_t, sigma_s0 = (
+                self.sigmas[self.step_index],
+                self.sigmas[self.step_index - 1],
+            )
+        else:
+            sigma_t, sigma_s0 = sigma, sigma_before
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - (i + 1)
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)
+        h_phi_k = h_phi_1 / hh - 1
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+        else:
+            D1s = None
+
+        if order == 1:
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+        else:
+            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            corr_res = (
+                torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+                if D1s is not None
+                else 0
+            )
+            D1_t = model_t - m0
+            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            corr_res = (
+                torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+                if D1s is not None
+                else 0
+            )
+            D1_t = model_t - m0
+            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+
+        return x_t.to(x.dtype)
+
+    def step_unipc(
+        self,
+        model_output,
+        timestep=None,
+        sample=None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', run 'set_timesteps' first"
+            )
+
+        if self.step_index is None:
+            self._step_index = 0
+
+        use_corrector = (
+            self.step_index > 0
+            and self.step_index - 1 not in self.disable_corrector
+            and self.last_sample is not None
+        )
+
+        model_output_convert = self.convert_model_output(model_output, sample=sample)
+
+        if use_corrector:
+            sample = self.multistep_uni_c_bh_update(
+                this_model_output=model_output_convert,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self.timestep_list[i] = self.timestep_list[i + 1]
+        self.model_outputs[-1] = model_output_convert
+        self.timestep_list[-1] = timestep
+
+        if self.config.lower_order_final:
+            this_order = min(
+                self.config.solver_order, len(self.timesteps) - self.step_index
+            )
+        else:
+            this_order = self.config.solver_order
+        self.this_order = min(this_order, self.lower_order_nums + 1)
+        assert self.this_order > 0
+
+        self.last_sample = sample
+        prev_sample = self.multistep_uni_p_bh_update(
+            model_output=model_output,
+            sample=sample,
+            order=self.this_order,
+        )
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return HeliosSchedulerOutput(prev_sample=prev_sample)
+
+    # ---------------------------------- DMD ----------------------------------
+    def add_noise(self, original_samples, noise, timestep, sigmas, timesteps):
+        sigmas = sigmas.to(noise.device)
+        timesteps = timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
+        )
+        sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps):
+        original_dtype = flow_pred.dtype
+        device = flow_pred.device
+        flow_pred, xt, sigmas, timesteps = (
+            x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps)
+        )
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
+        )
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)
+        x0_pred = xt - sigma_t * flow_pred
+        return x0_pred.to(original_dtype)
+
+    def step_dmd(
+        self,
+        model_output: torch.FloatTensor,
+        timestep=None,
+        sample: torch.FloatTensor = None,
+        return_dict: bool = True,
+        cur_sampling_step: int = 0,
+        dmd_noisy_tensor: torch.FloatTensor | None = None,
+        dmd_sigmas: torch.FloatTensor | None = None,
+        dmd_timesteps: torch.FloatTensor | None = None,
+        all_timesteps: torch.FloatTensor | None = None,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        pred_image_or_video = self.convert_flow_pred_to_x0(
+            flow_pred=model_output,
+            xt=sample,
+            timestep=torch.full(
+                (model_output.shape[0],),
+                timestep,
+                dtype=torch.long,
+                device=model_output.device,
+            ),
+            sigmas=dmd_sigmas,
+            timesteps=dmd_timesteps,
+        )
+        if cur_sampling_step < len(all_timesteps) - 1:
+            prev_sample = self.add_noise(
+                pred_image_or_video,
+                dmd_noisy_tensor,
+                torch.full(
+                    (model_output.shape[0],),
+                    all_timesteps[cur_sampling_step + 1],
+                    dtype=torch.long,
+                    device=model_output.device,
+                ),
+                sigmas=dmd_sigmas,
+                timesteps=dmd_timesteps,
+            )
+        else:
+            prev_sample = pred_image_or_video
+
+        if not return_dict:
+            return (prev_sample,)
+        return HeliosSchedulerOutput(prev_sample=prev_sample)
+
+    # ---------------------------------- Main step ----------------------------------
+    def step(
+        self,
+        model_output,
+        timestep=None,
+        sample=None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        if self.config.scheduler_type == "euler":
+            return self.step_euler(
+                model_output=model_output,
+                timestep=timestep,
+                sample=sample,
+                return_dict=return_dict,
+            )
+        elif self.config.scheduler_type == "unipc":
+            return self.step_unipc(
+                model_output=model_output,
+                timestep=timestep,
+                sample=sample,
+                return_dict=return_dict,
+            )
+        elif self.config.scheduler_type == "dmd":
+            return self.step_dmd(
+                model_output=model_output,
+                timestep=timestep,
+                sample=sample,
+                return_dict=return_dict,
+                **kwargs,
+            )
+        else:
+            raise NotImplementedError(
+                f"Scheduler type '{self.config.scheduler_type}' not implemented"
+            )
+
+    def reset_scheduler_history(self):
+        self.model_outputs = [None] * self.config.solver_order
+        self.timestep_list = [None] * self.config.solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = self.config.disable_corrector
+        self.solver_p = None
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+
+    def set_shift(self, shift: float):
+        """Update the shift parameter (called by SchedulerLoader after loading)."""
+        self.config.shift = shift
+        self.shift = shift
+
+    def __len__(self):
+        return self.config.num_train_timesteps
+
+
+# Alias for Helios-Distilled which uses "HeliosDMDScheduler" in scheduler_config.json
+HeliosDMDScheduler = HeliosScheduler
+
+EntryClass = [HeliosScheduler, "HeliosDMDScheduler"]

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

@@ -0,0 +1,142 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SelfForcingFlowMatchSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
+    config_name = "scheduler_config.json"
+    order = 1
+
+    @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,
+        *args,
+        **kwargs,
+    ):
+        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.set_timesteps(num_inference_steps)
+
+    def set_timesteps(
+        self,
+        num_inference_steps=100,
+        denoising_strength=1.0,
+        return_dict=False,
+        **kwargs,
+    ):
+        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])
+        self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: torch.FloatTensor,
+        sample: torch.FloatTensor,
+        to_final=False,
+        return_dict=False,
+        **kwargs,
+    ):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        elif timestep.ndim == 0:
+            # handles the case where timestep is a scalar, this occurs when we
+            # use this scheduler for ODE trajectory
+            timestep = timestep.unsqueeze(0)
+
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        timestep = timestep.to(model_output.device)
+
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
+        )
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
+            return (prev_sample,)
+        return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B*T, C, H, W]
+            - noise: the noise with shape [B*T, C, H, W]
+            - timestep: the timestep with shape [B*T]
+        Output: the corrupted latent with shape [B*T, C, H, W]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
+        )
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        return sample
+
+    def set_shift(self, shift: float) -> None:
+        self.shift = shift
+
+
+EntryClass = SelfForcingFlowMatchScheduler

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

@@ -0,0 +1,5 @@
+from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import (
+    LatentUpsampler,
+)
+
+__all__ = ["LatentUpsampler"]

python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py

@@ -0,0 +1,277 @@
+# Ported from https://github.com/Lightricks/LTX-2
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+
+
+class BlurDownsample(torch.nn.Module):
+    """Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel."""
+
+    def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None:
+        super().__init__()
+        assert dims in (2, 3)
+        assert isinstance(stride, int) and stride >= 1
+        assert kernel_size >= 3 and kernel_size % 2 == 1
+        self.dims = dims
+        self.stride = stride
+        self.kernel_size = kernel_size
+
+        k = torch.tensor([math.comb(kernel_size - 1, i) for i in range(kernel_size)])
+        k2d = k[:, None] @ k[None, :]
+        k2d = (k2d / k2d.sum()).float()
+        self.register_buffer("kernel", k2d[None, None, :, :])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.stride == 1:
+            return x
+        if self.dims == 2:
+            return self._apply_2d(x)
+        b, _, f, _, _ = x.shape
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+        x = self._apply_2d(x)
+        h2, w2 = x.shape[-2:]
+        x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2)
+        return x
+
+    def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor:
+        c = x2d.shape[1]
+        weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size)
+        x2d = F.conv2d(
+            x2d,
+            weight=weight,
+            bias=None,
+            stride=self.stride,
+            padding=self.kernel_size // 2,
+            groups=c,
+        )
+        return x2d
+
+
+class PixelShuffleND(torch.nn.Module):
+    """N-dimensional pixel shuffle for upsampling tensors."""
+
+    def __init__(self, dims: int, upscale_factors: Tuple[int, int, int] = (2, 2, 2)):
+        super().__init__()
+        assert dims in [1, 2, 3]
+        self.dims = dims
+        self.upscale_factors = upscale_factors
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.dims == 3:
+            return rearrange(
+                x,
+                "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+                p1=self.upscale_factors[0],
+                p2=self.upscale_factors[1],
+                p3=self.upscale_factors[2],
+            )
+        elif self.dims == 2:
+            return rearrange(
+                x,
+                "b (c p1 p2) h w -> b c (h p1) (w p2)",
+                p1=self.upscale_factors[0],
+                p2=self.upscale_factors[1],
+            )
+        elif self.dims == 1:
+            return rearrange(
+                x,
+                "b (c p1) f h w -> b c (f p1) h w",
+                p1=self.upscale_factors[0],
+            )
+        else:
+            raise ValueError(f"Unsupported dims: {self.dims}")
+
+
+class ResBlock(torch.nn.Module):
+    """Residual block with two conv layers, group norm, and SiLU activation."""
+
+    def __init__(
+        self, channels: int, mid_channels: Optional[int] = None, dims: int = 3
+    ):
+        super().__init__()
+        if mid_channels is None:
+            mid_channels = channels
+        conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
+        self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1)
+        self.norm1 = torch.nn.GroupNorm(32, mid_channels)
+        self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1)
+        self.norm2 = torch.nn.GroupNorm(32, channels)
+        self.activation = torch.nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.conv1(x)
+        # Fused GroupNorm + SiLU on the first norm of the block. The second
+        # norm (line below) is followed by `silu(norm + residual)`, which the
+        # current `apply_group_norm_silu` helper does not cover -- left on
+        # the eager path until a `group_norm_add_silu` helper exists.
+        x = apply_group_norm_silu(x, self.norm1, self.activation)
+        x = self.conv2(x)
+        x = self.norm2(x)
+        x = self.activation(x + residual)
+        return x
+
+
+def _rational_for_scale(scale: float) -> Tuple[int, int]:
+    mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)}
+    if float(scale) not in mapping:
+        raise ValueError(
+            f"Unsupported scale {scale}. Choose from {list(mapping.keys())}"
+        )
+    return mapping[float(scale)]
+
+
+class SpatialRationalResampler(torch.nn.Module):
+    """Fully-learned rational spatial scaling via PixelShuffle + anti-aliased downsample."""
+
+    def __init__(self, mid_channels: int, scale: float):
+        super().__init__()
+        self.scale = float(scale)
+        self.num, self.den = _rational_for_scale(self.scale)
+        self.conv = torch.nn.Conv2d(
+            mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1
+        )
+        self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
+        self.blur_down = BlurDownsample(dims=2, stride=self.den)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, _, f, _, _ = x.shape
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+        x = self.conv(x)
+        x = self.pixel_shuffle(x)
+        x = self.blur_down(x)
+        x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+        return x
+
+
+class LatentUpsampler(torch.nn.Module, LayerwiseOffloadableModuleMixin):
+    """
+    Upsample VAE latents spatially and/or temporally.
+
+    Args:
+        in_channels: Number of channels in the input latent.
+        mid_channels: Number of channels in the middle layers.
+        num_blocks_per_stage: Number of ResBlocks per stage (pre/post upsampling).
+        dims: Dimensionality of convolutions (2 or 3).
+        spatial_upsample: Whether to spatially upsample.
+        temporal_upsample: Whether to temporally upsample.
+        spatial_scale: Scale factor for spatial upsampling.
+        rational_resampler: Whether to use rational resampler for spatial upsampling.
+    """
+
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["res_blocks", "post_upsample_res_blocks"]
+
+    def __init__(
+        self,
+        in_channels: int = 128,
+        mid_channels: int = 512,
+        num_blocks_per_stage: int = 4,
+        dims: int = 3,
+        spatial_upsample: bool = True,
+        temporal_upsample: bool = False,
+        spatial_scale: float = 2.0,
+        rational_resampler: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.mid_channels = mid_channels
+        self.num_blocks_per_stage = num_blocks_per_stage
+        self.dims = dims
+        self.spatial_upsample = spatial_upsample
+        self.temporal_upsample = temporal_upsample
+        self.spatial_scale = float(spatial_scale)
+        self.rational_resampler = rational_resampler
+
+        conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
+
+        self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1)
+        self.initial_norm = torch.nn.GroupNorm(32, mid_channels)
+        self.initial_activation = torch.nn.SiLU()
+
+        self.res_blocks = torch.nn.ModuleList(
+            [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+        )
+
+        if spatial_upsample and temporal_upsample:
+            self.upsampler = torch.nn.Sequential(
+                torch.nn.Conv3d(
+                    mid_channels, 8 * mid_channels, kernel_size=3, padding=1
+                ),
+                PixelShuffleND(3),
+            )
+        elif spatial_upsample:
+            if rational_resampler:
+                self.upsampler = SpatialRationalResampler(
+                    mid_channels=mid_channels, scale=self.spatial_scale
+                )
+            else:
+                self.upsampler = torch.nn.Sequential(
+                    torch.nn.Conv2d(
+                        mid_channels, 4 * mid_channels, kernel_size=3, padding=1
+                    ),
+                    PixelShuffleND(2),
+                )
+        elif temporal_upsample:
+            self.upsampler = torch.nn.Sequential(
+                torch.nn.Conv3d(
+                    mid_channels, 2 * mid_channels, kernel_size=3, padding=1
+                ),
+                PixelShuffleND(1),
+            )
+        else:
+            raise ValueError(
+                "Either spatial_upsample or temporal_upsample must be True"
+            )
+
+        self.post_upsample_res_blocks = torch.nn.ModuleList(
+            [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+        )
+
+        self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1)
+
+    def forward(self, latent: torch.Tensor) -> torch.Tensor:
+        b, _, f, _, _ = latent.shape
+
+        if self.dims == 2:
+            x = rearrange(latent, "b c f h w -> (b f) c h w")
+            x = self.initial_conv(x)
+            x = apply_group_norm_silu(x, self.initial_norm, self.initial_activation)
+            for block in self.res_blocks:
+                x = block(x)
+            x = self.upsampler(x)
+            for block in self.post_upsample_res_blocks:
+                x = block(x)
+            x = self.final_conv(x)
+            x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+        else:
+            x = self.initial_conv(latent)
+            x = apply_group_norm_silu(x, self.initial_norm, self.initial_activation)
+            for block in self.res_blocks:
+                x = block(x)
+
+            if self.temporal_upsample:
+                x = self.upsampler(x)
+                x = x[:, :, 1:, :, :]
+            elif isinstance(self.upsampler, SpatialRationalResampler):
+                x = self.upsampler(x)
+            else:
+                x = rearrange(x, "b c f h w -> (b f) c h w")
+                x = self.upsampler(x)
+                x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+
+            for block in self.post_upsample_res_blocks:
+                x = block(x)
+            x = self.final_conv(x)
+
+        return x

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

@@ -0,0 +1,602 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+    FusedAttnProcessor2_0,
+)
+from diffusers.models.autoencoders.vae import (
+    Decoder,
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+    Encoder,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    can_install_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.parallel.diffusers_spatial import (
+    enable_diffusers_decoder_spatial_parallel,
+    spatial_parallel_diffusers_decode,
+)
+
+
+class AutoencoderKL(nn.Module, LayerwiseOffloadableModuleMixin):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
+        out_channels (int,  *optional*, defaults to 3): Number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+            Tuple of downsample block types.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+            Tuple of upsample block types.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
+            Tuple of block output channels.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
+        sample_size (`int`, *optional*, defaults to `32`): Sample input size.
+        scaling_factor (`float`, *optional*, defaults to 0.18215):
+            The component-wise standard deviation of the trained latent space computed using the first batch of the
+            training set. This is used to scale the latent space to have unit variance when training the diffusion
+            model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
+            diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
+            / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
+            Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper.
+        force_upcast (`bool`, *optional*, default to `True`):
+            If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
+            can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
+            can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+        mid_block_add_attention (`bool`, *optional*, default to `True`):
+            If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
+            mid_block will only have resnet blocks
+    """
+
+    layerwise_offload_dit_group_enabled = False
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
+    layer_names = ["encoder.down_blocks", "decoder.up_blocks"]
+
+    def __init__(
+        self,
+        config: FluxVAEConfig,
+    ):
+        super().__init__()
+        self.config = config
+        arch_config = config.arch_config
+
+        in_channels = arch_config.in_channels
+        out_channels = arch_config.out_channels
+        down_block_types = arch_config.down_block_types
+        up_block_types = arch_config.up_block_types
+        block_out_channels = arch_config.block_out_channels
+        layers_per_block = arch_config.layers_per_block
+        act_fn = arch_config.act_fn
+        latent_channels = arch_config.latent_channels
+        norm_num_groups = arch_config.norm_num_groups
+        sample_size = arch_config.sample_size
+        use_quant_conv = arch_config.use_quant_conv
+        use_post_quant_conv = arch_config.use_post_quant_conv
+        mid_block_add_attention = arch_config.mid_block_add_attention
+
+        # pass init params to Encoder
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        # pass init params to Decoder
+        self.decoder = Decoder(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.quant_conv = (
+            nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
+            if use_quant_conv
+            else None
+        )
+        self.post_quant_conv = (
+            nn.Conv2d(latent_channels, latent_channels, 1)
+            if use_post_quant_conv
+            else None
+        )
+
+        self.use_slicing = False
+        self.use_tiling = False
+        self.use_parallel_decode = config.use_parallel_decode
+        self.parallel_decode_mode = config.parallel_decode_mode
+        self._spatial_parallel_decode_enabled = False
+        self._spatial_parallel_upsample_count = 0
+        if can_install_spatial_shard_parallel_decode(self.config):
+            self._spatial_parallel_upsample_count = (
+                enable_diffusers_decoder_spatial_parallel(self.decoder)
+            )
+            self._spatial_parallel_decode_enabled = True
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = sample_size
+        sample_size = (
+            self.config.sample_size[0]
+            if isinstance(self.config.sample_size, (list, tuple))
+            else self.config.sample_size
+        )
+        self.tile_latent_min_size = int(
+            sample_size / (2 ** (len(self.config.block_out_channels) - 1))
+        )
+        self.tile_overlap_factor = 0.25
+
+    def enable_tiling(self, use_tiling: bool = True):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.use_tiling = use_tiling
+
+    def disable_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.enable_tiling(False)
+
+    def enable_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_slicing = True
+
+    def disable_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_slicing = False
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(
+            proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnAddedKVProcessor()
+        elif all(
+            proc.__class__ in CROSS_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, height, width = x.shape
+
+        if self.use_tiling and (
+            width > self.tile_sample_min_size or height > self.tile_sample_min_size
+        ):
+            return self._tiled_encode(x)
+
+        enc = self.encoder(x)
+        if self.quant_conv is not None:
+            enc = self.quant_conv(enc)
+
+        return enc
+
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_tiling and (
+            z.shape[-1] > self.tile_latent_min_size
+            or z.shape[-2] > self.tile_latent_min_size
+        ):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        if self.post_quant_conv is not None:
+            z = self.post_quant_conv(z)
+
+        if self._spatial_parallel_decode_enabled:
+            dec = spatial_parallel_diffusers_decode(
+                self.decoder, z, self._spatial_parallel_upsample_count
+            )
+        else:
+            dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def decode(self, z: torch.FloatTensor) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        return decoded
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[
+                :, :, y, :
+            ] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[
+                :, :, :, x
+            ] * (x / blend_extent)
+        return b
+
+    def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        enc = torch.cat(result_rows, dim=2)
+        return enc
+
+    def tiled_encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        moments = torch.cat(result_rows, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def tiled_decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[2], overlap_size):
+            row = []
+            for j in range(0, z.shape[3], overlap_size):
+                tile = z[
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                if self.config.use_post_quant_conv:
+                    tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        dec = torch.cat(result_rows, dim=2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        return dec
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        > [!WARNING] > This API is 🧪 experimental.
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError(
+                    "`fuse_qkv_projections()` is not supported for models having added KV projections."
+                )
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedAttnProcessor2_0())
+
+
+EntryClass = AutoencoderKL

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

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.parallel_conv import (
+    gather_and_trim_height,
+    split_height_for_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    can_install_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.parallel.diffusers_spatial import (
+    enable_diffusers_decoder_spatial_parallel,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class AutoencoderDC(nn.Module, LayerwiseOffloadableModuleMixin):
+    """Deep Compression Autoencoder wrapper with 32x spatial compression."""
+
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["_inner_model.encoder.down_blocks", "_inner_model.decoder.up_blocks"]
+
+    def __init__(self, config: SanaVAEConfig = None, **kwargs):
+        super().__init__()
+        self._config = config
+        self._inner_model = None
+        self._loaded_state_dict: dict[str, torch.Tensor] = {}
+        self._spatial_parallel_decode_enabled = False
+
+    def _ensure_inner_model(self, state_dict: dict[str, torch.Tensor] | None = None):
+        if self._inner_model is not None:
+            return
+
+        from diffusers import AutoencoderDC as DiffusersAutoencoderDC
+
+        device = "cpu"
+        state_to_load = (
+            state_dict if state_dict is not None else self._loaded_state_dict
+        )
+        if state_to_load:
+            first_tensor = next(iter(state_to_load.values()))
+            device = first_tensor.device
+        hf_config = {}
+        if self._config is not None:
+            arch = self._config.arch_config
+            for key, value in vars(arch).items():
+                if key == "extra_attrs" and isinstance(value, dict):
+                    for ek, ev in value.items():
+                        hf_config[ek] = ev
+                elif not key.startswith("_") and not callable(value):
+                    hf_config[key] = value
+
+        self._inner_model = DiffusersAutoencoderDC.from_config(hf_config)
+
+        if state_to_load:
+            missing, unexpected = self._inner_model.load_state_dict(
+                state_to_load, strict=False
+            )
+            if missing:
+                logger.warning(
+                    "AutoencoderDC missing keys when loading: %d keys", len(missing)
+                )
+                if len(missing) > 10:
+                    logger.debug("First 10 missing keys: %s", list(missing)[:10])
+                else:
+                    logger.debug("Missing keys: %s", list(missing))
+            if unexpected:
+                logger.debug(
+                    "AutoencoderDC unexpected keys when loading: %d keys",
+                    len(unexpected),
+                )
+            if state_dict is None:
+                self._loaded_state_dict.clear()
+
+        self._inner_model = self._inner_model.to(device)
+        if can_install_spatial_shard_parallel_decode(self._config):
+            enable_diffusers_decoder_spatial_parallel(self._inner_model.decoder)
+            self._spatial_parallel_decode_enabled = True
+
+    @property
+    def config(self):
+        if self._inner_model is not None:
+            return self._inner_model.config
+        return self._config
+
+    @property
+    def dtype(self):
+        if self._inner_model is not None:
+            return next(self._inner_model.parameters()).dtype
+        return torch.float32
+
+    @property
+    def device(self):
+        if self._inner_model is not None:
+            return next(self._inner_model.parameters()).device
+        return torch.device("cpu")
+
+    def encode(self, x: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        return self._inner_model.encode(x, **kwargs)
+
+    def decode(self, z: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        z = z.to(dtype=self.dtype)
+        if not self._spatial_parallel_decode_enabled:
+            return self._inner_model.decode(z, **kwargs)
+
+        expected_height = (
+            z.shape[-2] * self._config.arch_config.spatial_compression_ratio
+        )
+        z, expected_height = split_height_for_parallel_decode(
+            z,
+            expected_height=expected_height,
+            world_size=get_decode_parallel_world_size(),
+            rank=get_decode_parallel_rank(),
+        )
+        decoded = self._inner_model.decode(z, **kwargs)
+        if isinstance(decoded, tuple):
+            sample = gather_and_trim_height(decoded[0], expected_height)
+            return (sample, *decoded[1:])
+        sample = gather_and_trim_height(decoded.sample, expected_height)
+        return decoded.__class__(sample=sample)
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        return self._inner_model(x, **kwargs)
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, torch.Tensor],
+        strict: bool = True,
+        assign: bool = False,
+    ):
+        """Intercept load_state_dict to route weights into the inner diffusers model."""
+        self._ensure_inner_model(state_dict=state_dict)
+
+    def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]:
+        self._ensure_inner_model()
+        return self._inner_model.state_dict(*args, **kwargs)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Buffer weights for deferred loading. The inner model is built lazily."""
+        loaded_params: set[str] = set()
+        for name, weight in weights:
+            self._loaded_state_dict[name] = weight
+            loaded_params.add(name)
+        return loaded_params
+
+    def to(self, *args, **kwargs):
+        if self._inner_model is not None:
+            self._inner_model = self._inner_model.to(*args, **kwargs)
+        return super().to(*args, **kwargs)
+
+
+EntryClass = AutoencoderDC

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

@@ -0,0 +1,539 @@
+import math
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+)
+from diffusers.models.autoencoders.vae import (
+    Decoder,
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+    Encoder,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+
+from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    ParallelTiledVAE,
+    can_install_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.parallel.diffusers_spatial import (
+    enable_diffusers_decoder_spatial_parallel,
+    spatial_parallel_diffusers_decode,
+)
+
+
+class AutoencoderKLFlux2(ParallelTiledVAE):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
+
+    def __init__(
+        self,
+        config: Flux2VAEConfig,
+    ):
+        super().__init__(config=config)
+
+        self.config = config
+        arch_config = config.arch_config
+
+        in_channels: int = arch_config.in_channels
+        out_channels: int = arch_config.out_channels
+        down_block_types: Tuple[str, ...] = arch_config.down_block_types
+        up_block_types: Tuple[str, ...] = arch_config.up_block_types
+        block_out_channels: Tuple[int, ...] = arch_config.block_out_channels
+        decoder_block_out_channels: Optional[Tuple[int, ...]] = getattr(
+            arch_config, "decoder_block_out_channels", None
+        )
+        layers_per_block: int = arch_config.layers_per_block
+        act_fn: str = arch_config.act_fn
+        latent_channels: int = arch_config.latent_channels
+        norm_num_groups: int = arch_config.norm_num_groups
+        sample_size: int = arch_config.sample_size
+        use_quant_conv: bool = arch_config.use_quant_conv
+        use_post_quant_conv: bool = arch_config.use_post_quant_conv
+        mid_block_add_attention: bool = arch_config.mid_block_add_attention
+        batch_norm_eps: float = arch_config.batch_norm_eps
+        batch_norm_momentum: float = arch_config.batch_norm_momentum
+        patch_size: Tuple[int, int] = arch_config.patch_size
+        # pass init params to Encoder
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        # pass init params to Decoder
+        self.decoder = Decoder(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=decoder_block_out_channels or block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.quant_conv = (
+            nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
+            if use_quant_conv
+            else None
+        )
+        self.post_quant_conv = (
+            nn.Conv2d(latent_channels, latent_channels, 1)
+            if use_post_quant_conv
+            else None
+        )
+
+        self.bn = nn.BatchNorm2d(
+            math.prod(patch_size) * latent_channels,
+            eps=batch_norm_eps,
+            momentum=batch_norm_momentum,
+            affine=False,
+            track_running_stats=True,
+        )
+
+        self.use_slicing = False
+        self.use_tiling = False
+        self._spatial_parallel_decode_enabled = False
+        self._spatial_parallel_upsample_count = 0
+        if can_install_spatial_shard_parallel_decode(self.config):
+            self._spatial_parallel_upsample_count = (
+                enable_diffusers_decoder_spatial_parallel(self.decoder)
+            )
+            self._spatial_parallel_decode_enabled = True
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = self.config.sample_size
+        sample_size = (
+            self.config.sample_size[0]
+            if isinstance(self.config.sample_size, (list, tuple))
+            else self.config.sample_size
+        )
+        self.tile_latent_min_size = int(
+            sample_size / (2 ** (len(self.config.block_out_channels) - 1))
+        )
+        self.tile_overlap_factor = 0.25
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(
+            proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnAddedKVProcessor()
+        elif all(
+            proc.__class__ in CROSS_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, height, width = x.shape
+
+        if self.use_tiling and (
+            width > self.tile_sample_min_size or height > self.tile_sample_min_size
+        ):
+            return self._tiled_encode(x)
+
+        enc = self.encoder(x)
+        if self.quant_conv is not None:
+            enc = self.quant_conv(enc)
+
+        return enc
+
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[DiagonalGaussianDistribution]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+
+        if x.ndim == 5:
+            assert x.shape[2] == 1
+            x = x.squeeze(2)
+
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+        return posterior
+
+    def _decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_tiling and (
+            z.shape[-1] > self.tile_latent_min_size
+            or z.shape[-2] > self.tile_latent_min_size
+        ):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        if self.post_quant_conv is not None:
+            z = self.post_quant_conv(z)
+
+        if self._spatial_parallel_decode_enabled:
+            dec = spatial_parallel_diffusers_decode(
+                self.decoder, z, self._spatial_parallel_upsample_count
+            )
+        else:
+            dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def decode(
+        self, z: torch.FloatTensor, return_dict: bool = True, generator=None
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        return decoded
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[
+                :, :, y, :
+            ] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[
+                :, :, :, x
+            ] * (x / blend_extent)
+        return b
+
+    def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        enc = torch.cat(result_rows, dim=2)
+        return enc
+
+    def tiled_encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        moments = torch.cat(result_rows, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def tiled_decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[2], overlap_size):
+            row = []
+            for j in range(0, z.shape[3], overlap_size):
+                tile = z[
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                if self.config.use_post_quant_conv:
+                    tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        dec = torch.cat(result_rows, dim=2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+
+EntryClass = AutoencoderKLFlux2

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

@@ -0,0 +1,810 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from functools import lru_cache
+from math import isqrt, prod
+from typing import Optional, cast
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+
+from sglang.multimodal_gen.configs.models import VAEConfig
+from sglang.multimodal_gen.configs.models.vaes.base import (
+    should_use_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+    get_decode_parallel_group_coordinator,
+    get_decode_parallel_world_size,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+
+
+@lru_cache(maxsize=1)
+def _cached_decode_parallel_world_size(
+    is_dist_initialized: bool, is_model_parallel_initialized: bool, group_id: int
+) -> int:
+    if not is_dist_initialized or not is_model_parallel_initialized:
+        return 1
+    return get_decode_parallel_world_size()
+
+
+def _decode_parallel_world_size() -> int:
+    is_dist_initialized = dist.is_initialized()
+    is_model_parallel_initialized = model_parallel_is_initialized()
+    if not is_dist_initialized or not is_model_parallel_initialized:
+        return _cached_decode_parallel_world_size(
+            is_dist_initialized, is_model_parallel_initialized, 0
+        )
+    return _cached_decode_parallel_world_size(
+        is_dist_initialized,
+        is_model_parallel_initialized,
+        id(get_decode_parallel_group_coordinator()),
+    )
+
+
+def has_decode_parallel_world() -> bool:
+    return _decode_parallel_world_size() > 1
+
+
+def can_install_spatial_shard_parallel_decode(config: VAEConfig | None) -> bool:
+    world_size = _decode_parallel_world_size()
+    return (
+        config is not None
+        and world_size > 1
+        and should_use_spatial_shard_parallel_decode(config, world_size=world_size)
+    )
+
+
+def should_run_spatial_shard_parallel_decode(
+    config: VAEConfig, z: torch.Tensor
+) -> bool:
+    world_size = _decode_parallel_world_size()
+    return world_size > 1 and should_use_spatial_shard_parallel_decode(
+        config, z, world_size
+    )
+
+
+class ParallelTiledVAE(ABC, nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "encoder.down_blocks",
+        "decoder.up_blocks",
+        "encoder.down",
+        "decoder.up",
+    ]
+    tile_sample_min_height: int
+    tile_sample_min_width: int
+    tile_sample_min_num_frames: int
+    tile_sample_stride_height: int
+    tile_sample_stride_width: int
+    tile_sample_stride_num_frames: int
+    blend_num_frames: int
+    use_tiling: bool
+    use_temporal_tiling: bool
+    use_parallel_tiling: bool
+    use_parallel_decode: bool
+    parallel_decode_mode: str
+
+    def __init__(self, config: VAEConfig, **kwargs) -> None:
+        super().__init__()
+        self.config = config
+        self.tile_sample_min_height = config.tile_sample_min_height
+        self.tile_sample_min_width = config.tile_sample_min_width
+        self.tile_sample_min_num_frames = config.tile_sample_min_num_frames
+        self.tile_sample_stride_height = config.tile_sample_stride_height
+        self.tile_sample_stride_width = config.tile_sample_stride_width
+        self.tile_sample_stride_num_frames = config.tile_sample_stride_num_frames
+        self.blend_num_frames = config.blend_num_frames
+        self.use_tiling = config.use_tiling
+        self.use_temporal_tiling = config.use_temporal_tiling
+        self.use_parallel_tiling = config.use_parallel_tiling
+        self.use_parallel_decode = config.use_parallel_decode
+        self.parallel_decode_mode = config.parallel_decode_mode
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def temporal_compression_ratio(self) -> int:
+        return cast(int, self.config.temporal_compression_ratio)
+
+    @property
+    def spatial_compression_ratio(self) -> int:
+        return cast(int, self.config.spatial_compression_ratio)
+
+    @property
+    def scaling_factor(self) -> float | torch.Tensor:
+        return cast(float | torch.Tensor, self.config.scaling_factor)
+
+    @abstractmethod
+    def _encode(self, *args, **kwargs) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def _decode(self, *args, **kwargs) -> torch.Tensor:
+        pass
+
+    def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution:
+        batch_size, num_channels, num_frames, height, width = x.shape
+        latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
+
+        if (
+            self.use_tiling
+            and self.use_temporal_tiling
+            and num_frames > self.tile_sample_min_num_frames
+        ):
+            latents = self.tiled_encode(x)[:, :, :latent_num_frames]
+        elif self.use_tiling and (
+            width > self.tile_sample_min_width or height > self.tile_sample_min_height
+        ):
+            latents = self.spatial_tiled_encode(x)[:, :, :latent_num_frames]
+        else:
+            latents = self._encode(x)[:, :, :latent_num_frames]
+        return DiagonalGaussianDistribution(latents)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = z.shape
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+        tile_latent_min_num_frames = (
+            self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        )
+        num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
+
+        if should_run_spatial_shard_parallel_decode(self.config, z):
+            return self._decode(z)[:, :, :num_sample_frames]
+
+        if (
+            self.parallel_decode_mode == "tiled"
+            and self.use_tiling
+            and self.use_parallel_tiling
+            and get_sp_world_size() > 1
+        ):
+            return self.parallel_tiled_decode(z)[:, :, :num_sample_frames]
+        if (
+            self.use_tiling
+            and self.use_temporal_tiling
+            and num_frames > tile_latent_min_num_frames
+        ):
+            return self.tiled_decode(z)[:, :, :num_sample_frames]
+
+        if self.use_tiling and (
+            width > tile_latent_min_width or height > tile_latent_min_height
+        ):
+            return self.spatial_tiled_decode(z)[:, :, :num_sample_frames]
+
+        return self._decode(z)[:, :, :num_sample_frames]
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (
+                1 - y / blend_extent
+            ) + b[:, :, :, y, :] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (
+                1 - x / blend_extent
+            ) + b[:, :, :, :, x] * (x / blend_extent)
+        return b
+
+    def blend_t(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (
+                1 - x / blend_extent
+            ) + b[:, :, x, :, :] * (x / blend_extent)
+        return b
+
+    def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, _, height, width = x.shape
+        # latent_height = height // self.spatial_compression_ratio
+        # latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                tile = x[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_height,
+                    j : j + self.tile_sample_min_width,
+                ]
+                tile = self._encode(tile)
+                row.append(tile)
+            rows.append(row)
+
+        return self._merge_spatial_tiles(
+            rows,
+            blend_height,
+            blend_width,
+            tile_latent_stride_height,
+            tile_latent_stride_width,
+        )
+
+    def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
+        """
+        world_size, rank = get_sp_world_size(), get_sp_parallel_rank()
+        _, _, T, H, W = z.shape
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_min_num_frames = (
+            self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_num_frames = (
+            self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        )
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Calculate tile dimensions
+        num_t_tiles = (
+            T + tile_latent_stride_num_frames - 1
+        ) // tile_latent_stride_num_frames
+        num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height
+        num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width
+        total_spatial_tiles = num_h_tiles * num_w_tiles
+        total_tiles = num_t_tiles * total_spatial_tiles
+        tiles_per_rank = (total_tiles + world_size - 1) // world_size
+        start_tile_idx = rank * tiles_per_rank
+        end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles)
+
+        local_results = []
+        local_dim_metadata = []
+        for global_idx in range(start_tile_idx, end_tile_idx):
+            t_idx = global_idx // total_spatial_tiles
+            spatial_idx = global_idx % total_spatial_tiles
+            h_idx = spatial_idx // num_w_tiles
+            w_idx = spatial_idx % num_w_tiles
+
+            t_start = t_idx * tile_latent_stride_num_frames
+            h_start = h_idx * tile_latent_stride_height
+            w_start = w_idx * tile_latent_stride_width
+
+            tile = z[
+                :,
+                :,
+                t_start : t_start + tile_latent_min_num_frames + 1,
+                h_start : h_start + tile_latent_min_height,
+                w_start : w_start + tile_latent_min_width,
+            ]
+            decoded_tile = self._decode(tile)
+            if t_start > 0:
+                decoded_tile = decoded_tile[:, :, 1:, :, :]
+            local_results.append(decoded_tile.reshape(-1))
+            local_dim_metadata.append(decoded_tile.shape)
+
+        if local_results:
+            results = torch.cat(local_results, dim=0).contiguous()
+        else:
+            results = z.new_empty((0,), dtype=z.dtype)
+        del local_results
+
+        local_size = torch.tensor(
+            [results.size(0)], device=results.device, dtype=torch.int64
+        )
+        all_sizes = [
+            torch.zeros(1, device=results.device, dtype=torch.int64)
+            for _ in range(world_size)
+        ]
+        dist.all_gather(all_sizes, local_size)
+        max_size = max(size.item() for size in all_sizes)
+
+        padded_results = torch.zeros(
+            max_size, device=results.device, dtype=results.dtype
+        )
+        padded_results[: results.size(0)] = results
+
+        gathered_dim_metadata = [None] * world_size
+        gathered_results = (
+            torch.zeros_like(padded_results)
+            .repeat(world_size, *[1] * len(padded_results.shape))
+            .contiguous()
+        )
+        dist.all_gather_into_tensor(gathered_results, padded_results)
+        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
+        gathered_dim_metadata = cast(list[list[torch.Size]], gathered_dim_metadata)
+
+        data: list = [
+            [[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)]
+            for _ in range(num_t_tiles)
+        ]
+        global_idx = 0
+        for i, per_rank_metadata in enumerate(gathered_dim_metadata):
+            start_shape = 0
+            for shape in per_rank_metadata:
+                mul_shape = prod(shape)
+                current_data = gathered_results[
+                    i, start_shape : start_shape + mul_shape
+                ].reshape(shape)
+                t_idx = global_idx // total_spatial_tiles
+                spatial_idx = global_idx % total_spatial_tiles
+                h_idx = spatial_idx // num_w_tiles
+                w_idx = spatial_idx % num_w_tiles
+                data[t_idx][h_idx][w_idx] = current_data
+                start_shape += mul_shape
+                global_idx += 1
+
+        result_slices = []
+        last_slice_data = None
+        for i, tem_data in enumerate(data):
+            slice_data = self._merge_spatial_tiles(
+                tem_data,
+                blend_height,
+                blend_width,
+                self.tile_sample_stride_height,
+                self.tile_sample_stride_width,
+            )
+            if i > 0:
+                slice_data = self.blend_t(
+                    last_slice_data, slice_data, self.blend_num_frames
+                )
+                result_slices.append(
+                    slice_data[:, :, : self.tile_sample_stride_num_frames, :, :]
+                )
+            else:
+                result_slices.append(
+                    slice_data[:, :, : self.tile_sample_stride_num_frames + 1, :, :]
+                )
+            last_slice_data = slice_data
+        return torch.cat(result_slices, dim=2)
+
+    def parallel_patch_decode(self, z: torch.FloatTensor) -> torch.FloatTensor:
+        world_size, rank = get_sp_world_size(), get_sp_parallel_rank()
+        if world_size <= 1:
+            return self._decode(z)
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+        overlap_h = max(0, tile_latent_min_height - tile_latent_stride_height)
+        overlap_w = max(0, tile_latent_min_width - tile_latent_stride_width)
+        halo_h = overlap_h // 2
+        halo_w = overlap_w // 2
+
+        _, _, _, latent_h, latent_w = z.shape
+        scale = self.spatial_compression_ratio
+        out_h = latent_h * scale
+        out_w = latent_w * scale
+        root = isqrt(world_size)
+        grid_rows, grid_cols = 1, world_size
+        for rows in range(root, 0, -1):
+            if world_size % rows == 0:
+                grid_rows, grid_cols = rows, world_size // rows
+                break
+        patch_id = rank
+        patch_row = patch_id // grid_cols
+        patch_col = patch_id % grid_cols
+
+        h0 = (patch_row * latent_h) // grid_rows
+        h1 = ((patch_row + 1) * latent_h) // grid_rows
+        w0 = (patch_col * latent_w) // grid_cols
+        w1 = ((patch_col + 1) * latent_w) // grid_cols
+
+        ext_h0 = max(0, h0 - halo_h)
+        ext_h1 = min(latent_h, h1 + halo_h)
+        ext_w0 = max(0, w0 - halo_w)
+        ext_w1 = min(latent_w, w1 + halo_w)
+
+        local_patch = z[:, :, :, ext_h0:ext_h1, ext_w0:ext_w1]
+        decoded_patch = self._decode(local_patch)
+
+        crop_top = (h0 - ext_h0) * scale
+        crop_bottom = crop_top + (h1 - h0) * scale
+        crop_left = (w0 - ext_w0) * scale
+        crop_right = crop_left + (w1 - w0) * scale
+        decoded_core = decoded_patch[
+            :, :, :, crop_top:crop_bottom, crop_left:crop_right
+        ].contiguous()
+
+        local_result = decoded_core.reshape(-1)
+        local_dim_metadata = torch.tensor(
+            decoded_core.shape, device=z.device, dtype=torch.int64
+        )
+        local_position = torch.tensor(
+            [h0 * scale, h1 * scale, w0 * scale, w1 * scale],
+            device=z.device,
+            dtype=torch.int64,
+        )
+        gathered_positions = [
+            torch.empty_like(local_position) for _ in range(world_size)
+        ]
+        dist.all_gather(gathered_positions, local_position)
+
+        local_size = torch.tensor(
+            [local_result.size(0)], device=z.device, dtype=torch.int64
+        )
+        gathered_dim_metadata = [
+            torch.empty_like(local_dim_metadata) for _ in range(world_size)
+        ]
+        dist.all_gather(gathered_dim_metadata, local_dim_metadata)
+
+        all_sizes = [
+            torch.zeros(1, device=z.device, dtype=torch.int64)
+            for _ in range(world_size)
+        ]
+        dist.all_gather(all_sizes, local_size)
+        max_size = max(size.item() for size in all_sizes)
+
+        padded_results = torch.zeros(max_size, device=z.device, dtype=z.dtype)
+        padded_results[: local_result.size(0)] = local_result
+        gathered_results = torch.empty(
+            (world_size, *padded_results.shape),
+            device=padded_results.device,
+            dtype=padded_results.dtype,
+        )
+        dist.all_gather_into_tensor(gathered_results, padded_results)
+
+        dec = z.new_empty(
+            (
+                decoded_core.shape[0],
+                decoded_core.shape[1],
+                decoded_core.shape[2],
+                out_h,
+                out_w,
+            )
+        )
+        for src_rank, positions in enumerate(gathered_positions):
+            h_start, h_end, w_start, w_end = [int(x.item()) for x in positions]
+            shape = tuple(int(x.item()) for x in gathered_dim_metadata[src_rank])
+            patch = gathered_results[src_rank][: prod(shape)].reshape(shape)
+            dec[:, :, :, h_start:h_end, w_start:w_end] = patch
+        return dec
+
+    def _merge_spatial_tiles(
+        self, tiles, blend_height, blend_width, stride_height, stride_width
+    ) -> torch.Tensor:
+        """Helper function to merge spatial tiles with blending"""
+        result_rows = []
+        for i, row in enumerate(tiles):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(tiles[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :stride_height, :stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+
+        Returns:
+            `torch.Tensor`:
+                The decoded images.
+        """
+
+        _, _, _, height, width = z.shape
+        # sample_height = height * self.spatial_compression_ratio
+        # sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                tile = z[
+                    :,
+                    :,
+                    :,
+                    i : i + tile_latent_min_height,
+                    j : j + tile_latent_min_width,
+                ]
+                decoded = self._decode(tile)
+                row.append(decoded)
+            rows.append(row)
+        return self._merge_spatial_tiles(
+            rows,
+            blend_height,
+            blend_width,
+            self.tile_sample_stride_height,
+            self.tile_sample_stride_width,
+        )
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, num_frames, height, width = x.shape
+
+        # tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        tile_latent_stride_num_frames = (
+            self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        )
+
+        row = []
+        for i in range(0, num_frames, self.tile_sample_stride_num_frames):
+            tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :]
+            if self.use_tiling and (
+                height > self.tile_sample_min_height
+                or width > self.tile_sample_min_width
+            ):
+                tile = self.spatial_tiled_encode(tile)
+            else:
+                tile = self._encode(tile)
+            if i > 0:
+                tile = tile[:, :, 1:, :, :]
+            row.append(tile)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, self.blend_num_frames)
+                result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :])
+            else:
+                result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :])
+        enc = torch.cat(result_row, dim=2)
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = z.shape
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_min_num_frames = (
+            self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        )
+        tile_latent_stride_num_frames = (
+            self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        )
+
+        row = []
+        for i in range(0, num_frames, tile_latent_stride_num_frames):
+            tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :]
+            if self.use_tiling and (
+                tile.shape[-1] > tile_latent_min_width
+                or tile.shape[-2] > tile_latent_min_height
+            ):
+                decoded = self.spatial_tiled_decode(tile)
+            else:
+                decoded = self._decode(tile)
+            if i > 0:
+                decoded = decoded[:, :, 1:, :, :]
+            row.append(decoded)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, self.blend_num_frames)
+                result_row.append(
+                    tile[:, :, : self.tile_sample_stride_num_frames, :, :]
+                )
+            else:
+                result_row.append(
+                    tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :]
+                )
+
+        dec = torch.cat(result_row, dim=2)
+        return dec
+
+    def enable_tiling(
+        self,
+        tile_sample_min_height: int | None = None,
+        tile_sample_min_width: int | None = None,
+        tile_sample_min_num_frames: int | None = None,
+        tile_sample_stride_height: int | None = None,
+        tile_sample_stride_width: int | None = None,
+        tile_sample_stride_num_frames: int | None = None,
+        blend_num_frames: int | None = None,
+        use_tiling: bool | None = None,
+        use_temporal_tiling: bool | None = None,
+        use_parallel_tiling: bool | None = None,
+    ) -> None:
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+
+        Args:
+            tile_sample_min_height (`int`, *optional*):
+                The minimum height required for a sample to be separated into tiles across the height dimension.
+            tile_sample_min_width (`int`, *optional*):
+                The minimum width required for a sample to be separated into tiles across the width dimension.
+            tile_sample_min_num_frames (`int`, *optional*):
+                The minimum number of frames required for a sample to be separated into tiles across the frame
+                dimension.
+            tile_sample_stride_height (`int`, *optional*):
+                The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
+                no tiling artifacts produced across the height dimension.
+            tile_sample_stride_width (`int`, *optional*):
+                The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
+                artifacts produced across the width dimension.
+            tile_sample_stride_num_frames (`int`, *optional*):
+                The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts
+                produced across the frame dimension.
+        """
+        self.use_tiling = True
+        self.tile_sample_min_height = (
+            tile_sample_min_height or self.tile_sample_min_height
+        )
+        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+        self.tile_sample_min_num_frames = (
+            tile_sample_min_num_frames or self.tile_sample_min_num_frames
+        )
+        self.tile_sample_stride_height = (
+            tile_sample_stride_height or self.tile_sample_stride_height
+        )
+        self.tile_sample_stride_width = (
+            tile_sample_stride_width or self.tile_sample_stride_width
+        )
+        self.tile_sample_stride_num_frames = (
+            tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
+        )
+        if blend_num_frames is not None:
+            self.blend_num_frames = blend_num_frames
+        else:
+            self.blend_num_frames = (
+                self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+            )
+        self.use_tiling = use_tiling or self.use_tiling
+        self.use_temporal_tiling = use_temporal_tiling or self.use_temporal_tiling
+        self.use_parallel_tiling = use_parallel_tiling or self.use_parallel_tiling
+
+    def disable_tiling(self) -> None:
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_tiling = False
+
+
+# adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691
+class DiagonalGaussianDistribution:
+
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(
+                self.mean, device=self.parameters.device, dtype=self.parameters.dtype
+            )
+
+    def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = randn_tensor(
+            self.mean.shape,
+            generator=generator,
+            device=self.parameters.device,
+            dtype=self.parameters.dtype,
+        )
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(
+        self,
+        other: Optional["DiagonalGaussianDistribution"] = None,
+        dims: tuple[int, ...] = (1, 2, 3),
+    ) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=dims,
+                )
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var
+                    - 1.0
+                    - self.logvar
+                    + other.logvar,
+                    dim=dims,
+                )
+
+    def nll(
+        self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)
+    ) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self) -> torch.Tensor:
+        return self.mean

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

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Decoder-only audio tokenizer for the Cosmos3 sound modality."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+
+
+class Snake1d(nn.Module):
+    def __init__(self, hidden_dim: int, logscale: bool = True) -> None:
+        super().__init__()
+        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.logscale = logscale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        alpha = torch.exp(self.alpha) if self.logscale else self.alpha
+        beta = torch.exp(self.beta) if self.logscale else self.beta
+        x = x.reshape(shape[0], shape[1], -1)
+        x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+        return x.reshape(shape)
+
+
+class OobleckResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int = 1) -> None:
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.snake1 = Snake1d(dim)
+        self.conv1 = weight_norm(
+            nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad)
+        )
+        self.snake2 = Snake1d(dim)
+        self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.conv1(self.snake1(x))
+        y = self.conv2(self.snake2(y))
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class OobleckDecoderBlock(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        stride: int,
+        output_padding: int,
+    ) -> None:
+        super().__init__()
+        self.snake1 = Snake1d(input_dim)
+        self.conv_t1 = weight_norm(
+            nn.ConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=output_padding,
+            )
+        )
+        self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
+        self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
+        self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.snake1(x)
+        x = self.conv_t1(x)
+        x = self.res_unit1(x)
+        x = self.res_unit2(x)
+        return self.res_unit3(x)
+
+
+class OobleckDecoder(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        input_channels: int,
+        audio_channels: int,
+        upsampling_ratios: list[int],
+        channel_multiples: list[int],
+    ) -> None:
+        super().__init__()
+        strides = upsampling_ratios
+        mults = [1] + list(channel_multiples)
+
+        self.conv1 = weight_norm(
+            nn.Conv1d(input_channels, channels * mults[-1], kernel_size=7, padding=3)
+        )
+
+        blocks = []
+        for i, stride in enumerate(strides):
+            blocks.append(
+                OobleckDecoderBlock(
+                    input_dim=channels * mults[len(strides) - i],
+                    output_dim=channels * mults[len(strides) - i - 1],
+                    stride=stride,
+                    output_padding=stride % 2,
+                )
+            )
+        self.block = nn.ModuleList(blocks)
+        self.snake1 = Snake1d(channels)
+        self.conv2 = weight_norm(
+            nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        for layer in self.block:
+            x = layer(x)
+        x = self.snake1(x)
+        return self.conv2(x)
+
+
+def _cfg(config: dict[str, Any], *keys: str, default: Any = None) -> Any:
+    for k in keys:
+        v = config.get(k)
+        if v is not None:
+            return v
+    return default
+
+
+class Cosmos3AVAEAudioTokenizer(nn.Module):
+    """Cosmos3 audio tokenizer: latents → waveform via an Oobleck decoder stack."""
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        super().__init__()
+        self.sample_rate = int(
+            _cfg(config, "sampling_rate", "sample_rate", default=48000)
+        )
+        self.audio_channels = int(
+            _cfg(
+                config,
+                "dec_out_channels",
+                "audio_channels",
+                default=2 if bool(config.get("stereo", True)) else 1,
+            )
+        )
+        self.latent_channels = int(
+            _cfg(config, "vocoder_input_dim", "io_channels", "latent_ch", default=64)
+        )
+        dec_strides = [
+            int(s) for s in _cfg(config, "dec_strides", default=[2, 4, 5, 6, 8])
+        ]
+        self.hop_size = int(
+            _cfg(
+                config,
+                "hop_size",
+                default=math.prod(dec_strides) if dec_strides else 1920,
+            )
+        )
+        stride_product = math.prod(dec_strides)
+        if stride_product != self.hop_size:
+            raise ValueError(
+                "Cosmos3 AVAE dec_strides product must equal hop_size: "
+                f"product={stride_product}, hop_size={self.hop_size}."
+            )
+
+        norm = str(_cfg(config, "normalization_type", default="none"))
+        if bool(_cfg(config, "normalize_latents", default=False)) and norm == "none":
+            norm = "tanh"
+        self.normalization_type = norm
+        self.tanh_input_scale = float(_cfg(config, "tanh_input_scale", default=1.5))
+        self.tanh_output_scale = float(_cfg(config, "tanh_output_scale", default=3.5))
+        self.tanh_clamp = float(_cfg(config, "tanh_clamp", default=0.995))
+
+        self.decoder = OobleckDecoder(
+            channels=int(_cfg(config, "dec_dim", default=320)),
+            input_channels=self.latent_channels,
+            audio_channels=self.audio_channels,
+            upsampling_ratios=list(reversed(dec_strides)),
+            channel_multiples=list(
+                _cfg(config, "dec_c_mults", default=[1, 2, 4, 8, 16])
+            ),
+        )
+
+    @property
+    def temporal_compression_factor(self) -> int:
+        return self.hop_size
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        return int(num_audio_samples) // self.hop_size
+
+    def get_audio_num_samples(self, num_latent_samples: int) -> int:
+        return int(num_latent_samples) * self.hop_size
+
+    def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        if self.normalization_type == "tanh":
+            in_dtype = latent.dtype
+            x = torch.clamp(
+                latent.float() / self.tanh_output_scale,
+                -self.tanh_clamp,
+                self.tanh_clamp,
+            )
+            return (torch.atanh(x) * self.tanh_input_scale).to(in_dtype)
+        if self.normalization_type != "none":
+            raise ValueError(
+                f"Unsupported AVAE normalization_type={self.normalization_type!r}."
+            )
+        return latent
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        squeeze = latent.ndim == 2
+        if squeeze:
+            latent = latent.unsqueeze(0)
+        decoder_dtype = next(self.decoder.parameters()).dtype
+        decoder_device = next(self.decoder.parameters()).device
+        z = self._denormalize_latent(latent.to(decoder_device)).to(decoder_dtype)
+        audio = self.decoder(z).clamp(-1.0, 1.0).to(latent.dtype)
+        return audio.squeeze(0) if squeeze else audio
+
+
+EntryClass = Cosmos3AVAEAudioTokenizer

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

@@ -0,0 +1,646 @@
+# 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.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    DiagonalGaussianDistribution,
+)
+
+
+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
+
+
+# Scripting this brings model speed up 1.4x
+snake = torch.jit.script(_snake)
+
+
+# ROCm HIPRTC can fail to compile the scripted bf16 Snake kernel.
+def _should_use_eager_snake_on_rocm_bf16(x: torch.Tensor, alpha: torch.Tensor) -> bool:
+    return (
+        torch.version.hip is not None
+        and (x.is_cuda or alpha.is_cuda)
+        and (x.dtype == torch.bfloat16 or alpha.dtype == torch.bfloat16)
+    )
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        if _should_use_eager_snake_on_rocm_bf16(x, self.alpha):
+            return _snake(x, self.alpha)
+        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, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["encoder.block", "decoder.model"]
+
+    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/models/vaes/hunyuanvae.py

@@ -0,0 +1,978 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from diffusers
+
+# Copyright 2024 The Hunyuan Team, The HuggingFace Team and The sglang-diffusion Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu
+from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
+from sglang.multimodal_gen.runtime.layers.parallel_conv import (
+    SpatialParallelConv3d,
+    chunk_height_by_sizes,
+    disable_spatial_parallel_decode,
+    gather_and_trim_height,
+    gather_variable_height,
+    split_height_for_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    ParallelTiledVAE,
+    can_install_spatial_shard_parallel_decode,
+    should_run_spatial_shard_parallel_decode,
+)
+
+
+def prepare_causal_attention_mask(
+    num_frames: int,
+    height_width: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    batch_size: int | None = None,
+) -> torch.Tensor:
+    indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device)
+    indices_blocks = indices.repeat_interleave(height_width)
+    x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy")
+    mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype)
+
+    if batch_size is not None:
+        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+    return mask
+
+
+def _make_spatial_parallel_conv3d(
+    conv: nn.Conv3d,
+    *,
+    height_pad: int,
+    width_pad: int,
+    padding_mode: str,
+) -> SpatialParallelConv3d:
+    spatial_conv = SpatialParallelConv3d(
+        in_channels=conv.in_channels,
+        out_channels=conv.out_channels,
+        kernel_size=conv.kernel_size,
+        stride=conv.stride,
+        padding=(0, height_pad, width_pad),
+        dilation=conv.dilation,
+        groups=conv.groups,
+        bias=conv.bias is not None,
+        padding_mode=padding_mode,
+    )
+    spatial_conv.weight = conv.weight
+    spatial_conv.bias = conv.bias
+    return spatial_conv
+
+
+def _apply_group_norm_silu(
+    hidden_states: torch.Tensor,
+    norm: nn.GroupNorm,
+    nonlinearity: nn.Module,
+    spatial_parallel: bool,
+) -> torch.Tensor:
+    if not spatial_parallel:
+        return apply_group_norm_silu(hidden_states, norm, nonlinearity)
+    hidden_states, heights = gather_variable_height(hidden_states)
+    hidden_states = apply_group_norm_silu(
+        hidden_states.contiguous(), norm, nonlinearity
+    )
+    return chunk_height_by_sizes(hidden_states, heights)
+
+
+class HunyuanVAEAttention(nn.Module):
+
+    def __init__(
+        self, in_channels, heads, dim_head, eps, norm_num_groups, bias
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.heads = heads
+        self.dim_head = dim_head
+        self.eps = eps
+        self.norm_num_groups = norm_num_groups
+        self.bias = bias
+
+        inner_dim = heads * dim_head
+
+        # Define the projection layers
+        self.to_q = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_k = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_v = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, bias=bias))
+
+        # Optional normalization layers
+        self.group_norm = nn.GroupNorm(
+            norm_num_groups, in_channels, eps=eps, affine=True
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        # Project to query, key, value
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        # Reshape for multi-head attention
+        head_dim = self.dim_head
+
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+        # Perform scaled dot-product attention
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        # Reshape back
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, self.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        # Linear projection
+        hidden_states = self.to_out(hidden_states)
+
+        # Residual connection and rescale
+        hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
+class HunyuanVideoCausalConv3d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int, int] = 3,
+        stride: int | tuple[int, int, int] = 1,
+        padding: int | tuple[int, int, int] = 0,
+        dilation: int | tuple[int, int, int] = 1,
+        bias: bool = True,
+        pad_mode: str = "replicate",
+    ) -> None:
+        super().__init__()
+
+        kernel_size = (
+            (kernel_size, kernel_size, kernel_size)
+            if isinstance(kernel_size, int)
+            else kernel_size
+        )
+
+        self.pad_mode = pad_mode
+        self.time_causal_padding = (
+            kernel_size[0] // 2,
+            kernel_size[0] // 2,
+            kernel_size[1] // 2,
+            kernel_size[1] // 2,
+            kernel_size[2] - 1,
+            0,
+        )
+        self.spatial_parallel_time_padding = (
+            0,
+            0,
+            0,
+            0,
+            kernel_size[2] - 1,
+            0,
+        )
+        self.spatial_parallel = False
+
+        self.conv = nn.Conv3d(
+            in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias
+        )
+
+    def enable_spatial_parallel(self) -> None:
+        if isinstance(self.conv, SpatialParallelConv3d):
+            self.spatial_parallel = True
+            return
+        if self.conv.kernel_size == (1, 1, 1):
+            self.spatial_parallel = True
+            return
+        self.conv = _make_spatial_parallel_conv3d(
+            self.conv,
+            height_pad=self.time_causal_padding[2],
+            width_pad=self.time_causal_padding[0],
+            padding_mode=self.pad_mode,
+        )
+        self.spatial_parallel = True
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        padding = (
+            self.spatial_parallel_time_padding
+            if self.spatial_parallel
+            else self.time_causal_padding
+        )
+        if any(padding):
+            hidden_states = F.pad(hidden_states, padding, mode=self.pad_mode)
+        return self.conv(hidden_states)
+
+
+class HunyuanVideoUpsampleCausal3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        kernel_size: int = 3,
+        stride: int = 1,
+        bias: bool = True,
+        upsample_factor: tuple[int, ...] = (2, 2, 2),
+    ) -> None:
+        super().__init__()
+
+        out_channels = out_channels or in_channels
+        self.upsample_factor = upsample_factor
+
+        self.conv = HunyuanVideoCausalConv3d(
+            in_channels, out_channels, kernel_size, stride, bias=bias
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_frames = hidden_states.size(2)
+
+        first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2)
+        first_frame = F.interpolate(
+            first_frame.squeeze(2),
+            scale_factor=self.upsample_factor[1:],
+            mode="nearest",
+        ).unsqueeze(2)
+
+        if num_frames > 1:
+            # See: https://github.com/pytorch/pytorch/issues/81665
+            # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate
+            # is fixed, this will raise either a runtime error, or fail silently with bad outputs.
+            # If you are encountering an error here, make sure to try running encoding/decoding with
+            # `vae.enable_tiling()` first. If that doesn't work, open an issue at:
+            # https://github.com/huggingface/diffusers/issues
+            other_frames = other_frames.contiguous()
+            other_frames = F.interpolate(
+                other_frames, scale_factor=self.upsample_factor, mode="nearest"
+            )
+            hidden_states = torch.cat((first_frame, other_frames), dim=2)
+        else:
+            hidden_states = first_frame
+
+        hidden_states = self.conv(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideoDownsampleCausal3D(nn.Module):
+
+    def __init__(
+        self,
+        channels: int,
+        out_channels: int | None = None,
+        padding: int = 1,
+        kernel_size: int = 3,
+        bias: bool = True,
+        stride=2,
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or channels
+
+        self.conv = HunyuanVideoCausalConv3d(
+            channels, out_channels, kernel_size, stride, padding, bias=bias
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideoResnetBlockCausal3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        dropout: float = 0.0,
+        groups: int = 32,
+        eps: float = 1e-6,
+        non_linearity: str = "silu",
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+
+        self.nonlinearity = get_act_fn(non_linearity)
+
+        self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True)
+        self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0)
+
+        self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0)
+
+        self.conv_shortcut = None
+        if in_channels != out_channels:
+            self.conv_shortcut = HunyuanVideoCausalConv3d(
+                in_channels, out_channels, 1, 1, 0
+            )
+        self.spatial_parallel = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.contiguous()
+        residual = hidden_states
+
+        hidden_states = _apply_group_norm_silu(
+            hidden_states, self.norm1, self.nonlinearity, self.spatial_parallel
+        )
+        hidden_states = self.conv1(hidden_states)
+
+        hidden_states = _apply_group_norm_silu(
+            hidden_states, self.norm2, self.nonlinearity, self.spatial_parallel
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            residual = self.conv_shortcut(residual)
+
+        hidden_states = hidden_states + residual
+        return hidden_states
+
+
+class HunyuanVideoMidBlock3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+    ) -> None:
+        super().__init__()
+        resnet_groups = (
+            resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        )
+        self.add_attention = add_attention
+
+        # There is always at least one resnet
+        resnets = [
+            HunyuanVideoResnetBlockCausal3D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                non_linearity=resnet_act_fn,
+            )
+        ]
+        attentions: list[HunyuanVAEAttention | None] = []
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(
+                    HunyuanVAEAttention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        eps=resnet_eps,
+                        norm_num_groups=resnet_groups,
+                        bias=True,
+                    )
+                )
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        self.spatial_parallel = False
+        self.gradient_checkpointing = False
+
+    def _run_attention(
+        self, attn: HunyuanVAEAttention, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        heights = None
+        if self.spatial_parallel:
+            hidden_states, heights = gather_variable_height(hidden_states)
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3)
+        attention_mask = prepare_causal_attention_mask(
+            num_frames,
+            height * width,
+            hidden_states.dtype,
+            hidden_states.device,
+            batch_size=batch_size,
+        )
+        hidden_states = attn(hidden_states, attention_mask=attention_mask)
+        hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(
+            0, 4, 1, 2, 3
+        )
+        if heights is not None:
+            hidden_states = chunk_height_by_sizes(hidden_states, heights)
+        return hidden_states
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(
+                self.resnets[0], hidden_states
+            )
+
+            for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True):
+                if attn is not None:
+                    hidden_states = self._run_attention(attn, hidden_states)
+
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+
+        else:
+            hidden_states = self.resnets[0](hidden_states)
+
+            for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True):
+                if attn is not None:
+                    hidden_states = self._run_attention(attn, hidden_states)
+
+                hidden_states = resnet(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoDownBlock3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_downsample: bool = True,
+        downsample_stride: tuple[int, ...] | int = 2,
+        downsample_padding: int = 1,
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    HunyuanVideoDownsampleCausal3D(
+                        out_channels,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        stride=downsample_stride,
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for resnet in self.resnets:
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+        else:
+            for resnet in self.resnets:
+                hidden_states = resnet(hidden_states)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoUpBlock3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_upsample: bool = True,
+        upsample_scale_factor: tuple[int, ...] = (2, 2, 2),
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [
+                    HunyuanVideoUpsampleCausal3D(
+                        out_channels,
+                        out_channels=out_channels,
+                        upsample_factor=upsample_scale_factor,
+                    )
+                ]
+            )
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for resnet in self.resnets:
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+
+        else:
+            for resnet in self.resnets:
+                hidden_states = resnet(hidden_states)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoEncoder3D(nn.Module):
+    r"""
+    Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: tuple[str, ...] = (
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+        ),
+        block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        double_z: bool = True,
+        mid_block_add_attention=True,
+        temporal_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ) -> None:
+        super().__init__()
+
+        self.conv_in = HunyuanVideoCausalConv3d(
+            in_channels, block_out_channels[0], kernel_size=3, stride=1
+        )
+        self.mid_block: HunyuanVideoMidBlock3D | None = None
+        self.down_blocks = nn.ModuleList([])
+
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            if down_block_type != "HunyuanVideoDownBlock3D":
+                raise ValueError(f"Unsupported down_block_type: {down_block_type}")
+
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_downsample_layers = int(np.log2(temporal_compression_ratio))
+
+            if temporal_compression_ratio == 4:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(
+                    i >= (len(block_out_channels) - 1 - num_time_downsample_layers)
+                    and not is_final_block
+                )
+            elif temporal_compression_ratio == 8:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(i < num_time_downsample_layers)
+            else:
+                raise ValueError(
+                    f"Unsupported time_compression_ratio: {temporal_compression_ratio}"
+                )
+
+            downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
+            downsample_stride_T = (2,) if add_time_downsample else (1,)
+            downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
+
+            down_block = HunyuanVideoDownBlock3D(
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=bool(add_spatial_downsample or add_time_downsample),
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                downsample_stride=downsample_stride,
+                downsample_padding=0,
+            )
+
+            self.down_blocks.append(down_block)
+
+        self.mid_block = HunyuanVideoMidBlock3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            add_attention=mid_block_add_attention,
+        )
+
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6
+        )
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = HunyuanVideoCausalConv3d(
+            block_out_channels[-1], conv_out_channels, kernel_size=3
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for down_block in self.down_blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    down_block, hidden_states
+                )
+
+            hidden_states = self._gradient_checkpointing_func(
+                self.mid_block, hidden_states
+            )
+        else:
+            for down_block in self.down_blocks:
+                hidden_states = down_block(hidden_states)
+            assert self.mid_block is not None
+            hidden_states = self.mid_block(hidden_states)
+
+        hidden_states = apply_group_norm_silu(
+            hidden_states, self.conv_norm_out, self.conv_act
+        )
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoDecoder3D(nn.Module):
+    r"""
+    Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        up_block_types: tuple[str, ...] = (
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+        ),
+        block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        mid_block_add_attention=True,
+        time_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.conv_in = HunyuanVideoCausalConv3d(
+            in_channels, block_out_channels[-1], kernel_size=3, stride=1
+        )
+        self.up_blocks = nn.ModuleList([])
+
+        # mid
+        self.mid_block = HunyuanVideoMidBlock3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            add_attention=mid_block_add_attention,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            if up_block_type != "HunyuanVideoUpBlock3D":
+                raise ValueError(f"Unsupported up_block_type: {up_block_type}")
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_upsample_layers = int(np.log2(time_compression_ratio))
+
+            if time_compression_ratio == 4:
+                add_spatial_upsample = bool(i < num_spatial_upsample_layers)
+                add_time_upsample = bool(
+                    i >= len(block_out_channels) - 1 - num_time_upsample_layers
+                    and not is_final_block
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported time_compression_ratio: {time_compression_ratio}"
+                )
+
+            upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
+            upsample_scale_factor_T = (2,) if add_time_upsample else (1,)
+            upsample_scale_factor = tuple(
+                upsample_scale_factor_T + upsample_scale_factor_HW
+            )
+
+            up_block = HunyuanVideoUpBlock3D(
+                num_layers=self.layers_per_block + 1,
+                in_channels=prev_output_channel,
+                out_channels=output_channel,
+                add_upsample=bool(add_spatial_upsample or add_time_upsample),
+                upsample_scale_factor=upsample_scale_factor,
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+            )
+
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6
+        )
+        self.conv_act = nn.SiLU()
+        self.conv_out = HunyuanVideoCausalConv3d(
+            block_out_channels[0], out_channels, kernel_size=3
+        )
+
+        self.spatial_parallel = False
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(
+                self.mid_block, hidden_states
+            )
+
+            for up_block in self.up_blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    up_block, hidden_states
+                )
+        else:
+            hidden_states = self.mid_block(hidden_states)
+
+            for up_block in self.up_blocks:
+                hidden_states = up_block(hidden_states)
+
+        # post-process
+        hidden_states = _apply_group_norm_silu(
+            hidden_states,
+            self.conv_norm_out,
+            self.conv_act,
+            self.spatial_parallel,
+        )
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+def _enable_hunyuan_decoder_spatial_parallel(decoder: nn.Module) -> None:
+    for module in decoder.modules():
+        if isinstance(module, HunyuanVideoCausalConv3d):
+            module.enable_spatial_parallel()
+        elif isinstance(module, HunyuanVideoResnetBlockCausal3D):
+            module.spatial_parallel = True
+        elif isinstance(module, HunyuanVideoMidBlock3D):
+            module.spatial_parallel = True
+        elif isinstance(module, HunyuanVideoDecoder3D):
+            module.spatial_parallel = True
+
+
+class AutoencoderKLHunyuanVideo(ParallelTiledVAE):
+    r"""
+    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
+    Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603).
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        config: HunyuanVAEConfig,
+    ) -> None:
+        nn.Module.__init__(self)
+        ParallelTiledVAE.__init__(self, config)
+
+        # TODO(will): only pass in config. We do this by manually defining a
+        # config for hunyuan vae
+        self.block_out_channels = config.block_out_channels
+
+        if config.load_encoder:
+            self.encoder = HunyuanVideoEncoder3D(
+                in_channels=config.in_channels,
+                out_channels=config.latent_channels,
+                down_block_types=config.down_block_types,
+                block_out_channels=config.block_out_channels,
+                layers_per_block=config.layers_per_block,
+                norm_num_groups=config.norm_num_groups,
+                act_fn=config.act_fn,
+                double_z=True,
+                mid_block_add_attention=config.mid_block_add_attention,
+                temporal_compression_ratio=config.temporal_compression_ratio,
+                spatial_compression_ratio=config.spatial_compression_ratio,
+            )
+            self.quant_conv = nn.Conv3d(
+                2 * config.latent_channels, 2 * config.latent_channels, kernel_size=1
+            )
+
+        if config.load_decoder:
+            self.decoder = HunyuanVideoDecoder3D(
+                in_channels=config.latent_channels,
+                out_channels=config.out_channels,
+                up_block_types=config.up_block_types,
+                block_out_channels=config.block_out_channels,
+                layers_per_block=config.layers_per_block,
+                norm_num_groups=config.norm_num_groups,
+                act_fn=config.act_fn,
+                time_compression_ratio=config.temporal_compression_ratio,
+                spatial_compression_ratio=config.spatial_compression_ratio,
+                mid_block_add_attention=config.mid_block_add_attention,
+            )
+            self.post_quant_conv = nn.Conv3d(
+                config.latent_channels, config.latent_channels, kernel_size=1
+            )
+            self._spatial_parallel_decode_enabled = False
+            if can_install_spatial_shard_parallel_decode(self.config):
+                _enable_hunyuan_decoder_spatial_parallel(self.decoder)
+                self._spatial_parallel_decode_enabled = True
+
+    def _should_use_spatial_parallel_decode(self, z: torch.Tensor) -> bool:
+        return (
+            self._spatial_parallel_decode_enabled
+            and should_run_spatial_shard_parallel_decode(self.config, z)
+        )
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.encoder(x)
+        enc = self.quant_conv(x)
+        return enc
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        z = self.post_quant_conv(z)
+        if self._should_use_spatial_parallel_decode(z):
+            expected_height = (
+                z.shape[-2] * self.config.arch_config.spatial_compression_ratio
+            )
+            z, expected_height = split_height_for_parallel_decode(
+                z,
+                expected_height=expected_height,
+                world_size=get_decode_parallel_world_size(),
+                rank=get_decode_parallel_rank(),
+            )
+            dec = gather_and_trim_height(self.decoder(z), expected_height)
+        elif self._spatial_parallel_decode_enabled:
+            with disable_spatial_parallel_decode():
+                dec = self.decoder(z)
+        else:
+            dec = self.decoder(z)
+        return dec
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec
+
+
+EntryClass = AutoencoderKLHunyuanVideo

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

@@ -0,0 +1,210 @@
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.ltx_2_vae import (
+    LTX2VideoCausalConv3d,
+    LTX2VideoResnetBlock3d,
+    LTXVideoDownsampler3d,
+)
+
+
+def _patchify_video(sample: torch.Tensor, patch_size: int) -> torch.Tensor:
+    if patch_size == 1:
+        return sample
+    batch_size, channels, num_frames, height, width = sample.shape
+    sample = sample.reshape(
+        batch_size,
+        channels,
+        num_frames,
+        1,
+        height // patch_size,
+        patch_size,
+        width // patch_size,
+        patch_size,
+    )
+    return sample.permute(0, 1, 3, 7, 5, 2, 4, 6).flatten(1, 4)
+
+
+class LTX23VideoPixelNorm(nn.Module):
+    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
+        rms = torch.sqrt(mean_sq + self.eps)
+        return x / rms
+
+
+class LTX23PerChannelStatistics(nn.Module):
+    def __init__(self, latent_channels: int) -> None:
+        super().__init__()
+        self.register_buffer("std-of-means", torch.empty(latent_channels))
+        self.register_buffer("mean-of-means", torch.empty(latent_channels))
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)
+        std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)
+        return (x - mean) / std
+
+
+class LTX23VideoResBlockStack(nn.Module):
+    def __init__(
+        self, channels: int, num_layers: int, spatial_padding_mode: str
+    ) -> None:
+        super().__init__()
+        self.res_blocks = nn.ModuleList(
+            [
+                LTX2VideoResnetBlock3d(
+                    in_channels=channels,
+                    out_channels=channels,
+                    spatial_padding_mode=spatial_padding_mode,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for res_block in self.res_blocks:
+            hidden_states = res_block(hidden_states, causal=True)
+        return hidden_states
+
+
+def _make_ltx23_encoder_block(
+    block_name: str,
+    block_config: dict[str, Any],
+    in_channels: int,
+    spatial_padding_mode: str,
+) -> tuple[nn.Module, int]:
+    if block_name == "res_x":
+        return (
+            LTX23VideoResBlockStack(
+                channels=in_channels,
+                num_layers=int(block_config["num_layers"]),
+                spatial_padding_mode=spatial_padding_mode,
+            ),
+            in_channels,
+        )
+
+    multiplier = int(block_config.get("multiplier", 2))
+    stride_map = {
+        "compress_space_res": (1, 2, 2),
+        "compress_time_res": (2, 1, 1),
+        "compress_all_res": (2, 2, 2),
+    }
+    stride = stride_map.get(block_name)
+    if stride is None:
+        raise ValueError(f"Unsupported LTX-2.3 encoder block: {block_name}")
+    out_channels = in_channels * multiplier
+    return (
+        LTXVideoDownsampler3d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            stride=stride,
+            spatial_padding_mode=spatial_padding_mode,
+        ),
+        out_channels,
+    )
+
+
+class LTX23VideoConditionEncoder(nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["down_blocks"]
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        super().__init__()
+
+        vae_config = config.get("vae", config)
+        latent_channels = int(vae_config["latent_channels"])
+        patch_size = int(vae_config.get("patch_size", 4))
+        spatial_padding_mode = str(vae_config.get("spatial_padding_mode", "zeros"))
+        encoder_blocks = list(vae_config["encoder_blocks"])
+        latent_log_var = str(vae_config.get("latent_log_var", "uniform"))
+
+        self.patch_size = patch_size
+        self.latency_channels = latent_channels
+        self.latent_log_var = latent_log_var
+        self.per_channel_statistics = LTX23PerChannelStatistics(latent_channels)
+
+        feature_channels = latent_channels
+        self.conv_in = LTX2VideoCausalConv3d(
+            in_channels=int(vae_config.get("in_channels", 3)) * patch_size**2,
+            out_channels=feature_channels,
+            kernel_size=3,
+            stride=1,
+            spatial_padding_mode=spatial_padding_mode,
+        )
+
+        self.down_blocks = nn.ModuleList()
+        for block_name, block_params in encoder_blocks:
+            block_config = (
+                {"num_layers": block_params}
+                if isinstance(block_params, int)
+                else dict(block_params)
+            )
+            block, feature_channels = _make_ltx23_encoder_block(
+                block_name=block_name,
+                block_config=block_config,
+                in_channels=feature_channels,
+                spatial_padding_mode=spatial_padding_mode,
+            )
+            self.down_blocks.append(block)
+
+        self.conv_norm_out = LTX23VideoPixelNorm(dim=1, eps=1e-8)
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = latent_channels
+        if latent_log_var == "per_channel":
+            conv_out_channels *= 2
+        elif latent_log_var in {"uniform", "constant"}:
+            conv_out_channels += 1
+        elif latent_log_var != "none":
+            raise ValueError(f"Unsupported latent_log_var: {latent_log_var}")
+
+        self.conv_out = LTX2VideoCausalConv3d(
+            in_channels=feature_channels,
+            out_channels=conv_out_channels,
+            kernel_size=3,
+            stride=1,
+            spatial_padding_mode=spatial_padding_mode,
+        )
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        frames_count = int(sample.shape[2])
+        if (frames_count - 1) % 8 != 0:
+            frames_to_crop = (frames_count - 1) % 8
+            sample = sample[:, :, :-frames_to_crop, ...]
+
+        hidden_states = _patchify_video(sample, self.patch_size)
+        hidden_states = self.conv_in(hidden_states, causal=True)
+
+        for block in self.down_blocks:
+            hidden_states = block(hidden_states)
+
+        hidden_states = self.conv_norm_out(hidden_states)
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states, causal=True)
+
+        if self.latent_log_var == "uniform":
+            means = hidden_states[:, :-1, ...]
+            logvar = hidden_states[:, -1:, ...]
+            hidden_states = torch.cat(
+                [
+                    means,
+                    logvar.repeat(1, means.shape[1], *([1] * (means.ndim - 2))),
+                ],
+                dim=1,
+            )
+        elif self.latent_log_var == "constant":
+            means = hidden_states[:, :-1, ...]
+            logvar = torch.full_like(means, -30.0)
+            hidden_states = torch.cat([means, logvar], dim=1)
+
+        means, _ = torch.chunk(hidden_states, 2, dim=1)
+        return self.per_channel_statistics.normalize(means)

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

@@ -0,0 +1,918 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders.vae import (
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.ltx_audio import LTXAudioVAEConfig
+from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
+
+LATENT_DOWNSAMPLE_FACTOR = 4
+
+
+class LTX2AudioCausalConv2d(nn.Module):
+    """
+    A causal 2D convolution that pads asymmetrically along the causal axis.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int]],
+        stride: int = 1,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causality_axis: str = "height",
+    ) -> None:
+        super().__init__()
+
+        self.causality_axis = causality_axis
+        kernel_size = (
+            (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+        )
+        dilation = (dilation, dilation) if isinstance(dilation, int) else dilation
+
+        pad_h = (kernel_size[0] - 1) * dilation[0]
+        pad_w = (kernel_size[1] - 1) * dilation[1]
+
+        if self.causality_axis == "none":
+            padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
+        elif self.causality_axis in {"width", "width-compatibility"}:
+            padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
+        elif self.causality_axis == "height":
+            padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
+        else:
+            raise ValueError(f"Invalid causality_axis: {causality_axis}")
+
+        self.padding = padding
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+
+class LTX2AudioPixelNorm(nn.Module):
+    """
+    Per-pixel (per-location) RMS normalization layer.
+    """
+
+    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
+        rms = torch.sqrt(mean_sq + self.eps)
+        return x / rms
+
+
+class LTX2AudioAttnBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        norm_type: str = "group",
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+
+        if norm_type == "group":
+            self.norm = nn.GroupNorm(
+                num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = self.norm(x)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        batch, channels, height, width = q.shape
+        q = q.reshape(batch, channels, height * width).permute(0, 2, 1).contiguous()
+        k = k.reshape(batch, channels, height * width).contiguous()
+        attn = torch.bmm(q, k) * (int(channels) ** (-0.5))
+        attn = torch.nn.functional.softmax(attn, dim=2)
+
+        v = v.reshape(batch, channels, height * width)
+        attn = attn.permute(0, 2, 1).contiguous()
+        h_ = torch.bmm(v, attn).reshape(batch, channels, height, width)
+
+        h_ = self.proj_out(h_)
+        return x + h_
+
+
+class LTX2AudioResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 512,
+        norm_type: str = "group",
+        causality_axis: str = "height",
+    ) -> None:
+        super().__init__()
+        self.causality_axis = causality_axis
+
+        if (
+            self.causality_axis is not None
+            and self.causality_axis != "none"
+            and norm_type == "group"
+        ):
+            raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        if norm_type == "group":
+            self.norm1 = nn.GroupNorm(
+                num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm1 = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.non_linearity = nn.SiLU()
+        if causality_axis is not None:
+            self.conv1 = LTX2AudioCausalConv2d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=causality_axis,
+            )
+        else:
+            self.conv1 = nn.Conv2d(
+                in_channels, out_channels, kernel_size=3, stride=1, padding=1
+            )
+        if temb_channels > 0:
+            self.temb_proj = nn.Linear(temb_channels, out_channels)
+        if norm_type == "group":
+            self.norm2 = nn.GroupNorm(
+                num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm2 = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.dropout = nn.Dropout(dropout)
+        if causality_axis is not None:
+            self.conv2 = LTX2AudioCausalConv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=causality_axis,
+            )
+        else:
+            self.conv2 = nn.Conv2d(
+                out_channels, out_channels, kernel_size=3, stride=1, padding=1
+            )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                if causality_axis is not None:
+                    self.conv_shortcut = LTX2AudioCausalConv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=3,
+                        stride=1,
+                        causality_axis=causality_axis,
+                    )
+                else:
+                    self.conv_shortcut = nn.Conv2d(
+                        in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                    )
+            else:
+                if causality_axis is not None:
+                    self.nin_shortcut = LTX2AudioCausalConv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        causality_axis=causality_axis,
+                    )
+                else:
+                    self.nin_shortcut = nn.Conv2d(
+                        in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                    )
+
+    def forward(
+        self, x: torch.Tensor, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = self.norm1(x)
+        h = self.non_linearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = self.non_linearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = (
+                self.conv_shortcut(x)
+                if self.use_conv_shortcut
+                else self.nin_shortcut(x)
+            )
+
+        return x + h
+
+
+class LTX2AudioDownsample(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: Optional[str] = "height",
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            # Padding tuple is in the order: (left, right, top, bottom).
+            if self.causality_axis == "none":
+                pad = (0, 1, 0, 1)
+            elif self.causality_axis == "width":
+                pad = (2, 0, 0, 1)
+            elif self.causality_axis == "height":
+                pad = (0, 1, 2, 0)
+            elif self.causality_axis == "width-compatibility":
+                pad = (1, 0, 0, 1)
+            else:
+                raise ValueError(
+                    f"Invalid `causality_axis` {self.causality_axis}; supported values are `none`, `width`, `height`,"
+                    f" and `width-compatibility`."
+                )
+
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            # with_conv=False implies that causality_axis is "none"
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class LTX2AudioUpsample(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: Optional[str] = "height",
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+        if self.with_conv:
+            if causality_axis is not None:
+                self.conv = LTX2AudioCausalConv2d(
+                    in_channels,
+                    in_channels,
+                    kernel_size=3,
+                    stride=1,
+                    causality_axis=causality_axis,
+                )
+            else:
+                self.conv = nn.Conv2d(
+                    in_channels, in_channels, kernel_size=3, stride=1, padding=1
+                )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+            if self.causality_axis is None or self.causality_axis == "none":
+                pass
+            elif self.causality_axis == "height":
+                x = x[:, :, 1:, :]
+            elif self.causality_axis == "width":
+                x = x[:, :, :, 1:]
+            elif self.causality_axis == "width-compatibility":
+                pass
+            else:
+                raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
+
+        return x
+
+
+class LTX2AudioAudioPatchifier:
+    """
+    Patchifier for spectrogram/audio latents.
+    """
+
+    def __init__(
+        self,
+        patch_size: int,
+        sample_rate: int = 16000,
+        hop_length: int = 160,
+        audio_latent_downsample_factor: int = 4,
+        is_causal: bool = True,
+    ):
+        self.hop_length = hop_length
+        self.sample_rate = sample_rate
+        self.audio_latent_downsample_factor = audio_latent_downsample_factor
+        self.is_causal = is_causal
+        self._patch_size = (1, patch_size, patch_size)
+
+    def patchify(self, audio_latents: torch.Tensor) -> torch.Tensor:
+        batch, channels, time, freq = audio_latents.shape
+        return audio_latents.permute(0, 2, 1, 3).reshape(batch, time, channels * freq)
+
+    def unpatchify(
+        self, audio_latents: torch.Tensor, channels: int, mel_bins: int
+    ) -> torch.Tensor:
+        batch, time, _ = audio_latents.shape
+        return audio_latents.view(batch, time, channels, mel_bins).permute(0, 2, 1, 3)
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        return self._patch_size
+
+
+class LTX2AudioEncoder(nn.Module):
+    def __init__(
+        self,
+        base_channels: int = 128,
+        output_channels: int = 1,
+        num_res_blocks: int = 2,
+        attn_resolutions: Optional[Tuple[int, ...]] = None,
+        in_channels: int = 2,
+        resolution: int = 256,
+        latent_channels: int = 8,
+        ch_mult: Tuple[int, ...] = (1, 2, 4),
+        norm_type: str = "group",
+        causality_axis: Optional[str] = "width",
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = False,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: Optional[int] = 64,
+        double_z: bool = True,
+    ):
+        super().__init__()
+
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+
+        self.base_channels = base_channels
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_ch = output_channels
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.latent_channels = latent_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+
+        base_block_channels = base_channels
+        base_resolution = resolution
+        self.z_shape = (1, latent_channels, base_resolution, base_resolution)
+
+        if self.causality_axis is not None:
+            self.conv_in = LTX2AudioCausalConv2d(
+                in_channels,
+                base_block_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_in = nn.Conv2d(
+                in_channels, base_block_channels, kernel_size=3, stride=1, padding=1
+            )
+
+        self.down = nn.ModuleList()
+        block_in = base_block_channels
+        curr_res = self.resolution
+
+        for level in range(self.num_resolutions):
+            stage = nn.Module()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_out = self.base_channels * self.channel_multipliers[level]
+
+            for _ in range(self.num_res_blocks):
+                stage.block.append(
+                    LTX2AudioResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                        norm_type=self.norm_type,
+                        causality_axis=self.causality_axis,
+                    )
+                )
+                block_in = block_out
+                if self.attn_resolutions:
+                    if curr_res in self.attn_resolutions:
+                        stage.attn.append(
+                            LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+                        )
+
+            if level != self.num_resolutions - 1:
+                stage.downsample = LTX2AudioDownsample(
+                    block_in, True, causality_axis=self.causality_axis
+                )
+                curr_res = curr_res // 2
+
+            self.down.append(stage)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = LTX2AudioResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+        if mid_block_add_attention:
+            self.mid.attn_1 = LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+        else:
+            self.mid.attn_1 = nn.Identity()
+        self.mid.block_2 = LTX2AudioResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+
+        final_block_channels = block_in
+        z_channels = 2 * latent_channels if double_z else latent_channels
+        if self.norm_type == "group":
+            self.norm_out = nn.GroupNorm(
+                num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True
+            )
+        elif self.norm_type == "pixel":
+            self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {self.norm_type}")
+        self.non_linearity = nn.SiLU()
+
+        if self.causality_axis is not None:
+            self.conv_out = LTX2AudioCausalConv2d(
+                final_block_channels,
+                z_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_out = nn.Conv2d(
+                final_block_channels, z_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states expected shape: (batch_size, channels, time, num_mel_bins)
+        hidden_states = self.conv_in(hidden_states)
+
+        for level in range(self.num_resolutions):
+            stage = self.down[level]
+            for block_idx, block in enumerate(stage.block):
+                hidden_states = block(hidden_states, temb=None)
+                if stage.attn:
+                    hidden_states = stage.attn[block_idx](hidden_states)
+
+            if level != self.num_resolutions - 1 and hasattr(stage, "downsample"):
+                hidden_states = stage.downsample(hidden_states)
+
+        hidden_states = self.mid.block_1(hidden_states, temb=None)
+        hidden_states = self.mid.attn_1(hidden_states)
+        hidden_states = self.mid.block_2(hidden_states, temb=None)
+
+        hidden_states = self.norm_out(hidden_states)
+        hidden_states = self.non_linearity(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+class LTX2AudioDecoder(nn.Module):
+    """
+    Symmetric decoder that reconstructs audio spectrograms from latent features.
+
+    The decoder mirrors the encoder structure with configurable channel multipliers, attention resolutions, and causal
+    convolutions.
+    """
+
+    def __init__(
+        self,
+        base_channels: int = 128,
+        output_channels: int = 1,
+        num_res_blocks: int = 2,
+        attn_resolutions: Optional[Tuple[int, ...]] = None,
+        in_channels: int = 2,
+        resolution: int = 256,
+        latent_channels: int = 8,
+        ch_mult: Tuple[int, ...] = (1, 2, 4),
+        norm_type: str = "group",
+        causality_axis: Optional[str] = "width",
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = False,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: Optional[int] = 64,
+    ) -> None:
+        super().__init__()
+
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+        self.patchifier = LTX2AudioAudioPatchifier(
+            patch_size=1,
+            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+            sample_rate=sample_rate,
+            hop_length=mel_hop_length,
+            is_causal=is_causal,
+        )
+
+        self.base_channels = base_channels
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_ch = output_channels
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.latent_channels = latent_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+
+        base_block_channels = base_channels * self.channel_multipliers[-1]
+        base_resolution = resolution // (2 ** (self.num_resolutions - 1))
+        self.z_shape = (1, latent_channels, base_resolution, base_resolution)
+
+        if self.causality_axis is not None:
+            self.conv_in = LTX2AudioCausalConv2d(
+                latent_channels,
+                base_block_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_in = nn.Conv2d(
+                latent_channels, base_block_channels, kernel_size=3, stride=1, padding=1
+            )
+        self.non_linearity = nn.SiLU()
+        self.mid = nn.Module()
+        self.mid.block_1 = LTX2AudioResnetBlock(
+            in_channels=base_block_channels,
+            out_channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+        if mid_block_add_attention:
+            self.mid.attn_1 = LTX2AudioAttnBlock(
+                base_block_channels, norm_type=self.norm_type
+            )
+        else:
+            self.mid.attn_1 = nn.Identity()
+        self.mid.block_2 = LTX2AudioResnetBlock(
+            in_channels=base_block_channels,
+            out_channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+
+        self.up = nn.ModuleList()
+        block_in = base_block_channels
+        curr_res = self.resolution // (2 ** (self.num_resolutions - 1))
+
+        for level in reversed(range(self.num_resolutions)):
+            stage = nn.Module()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_out = self.base_channels * self.channel_multipliers[level]
+
+            for _ in range(self.num_res_blocks + 1):
+                stage.block.append(
+                    LTX2AudioResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                        norm_type=self.norm_type,
+                        causality_axis=self.causality_axis,
+                    )
+                )
+                block_in = block_out
+                if self.attn_resolutions:
+                    if curr_res in self.attn_resolutions:
+                        stage.attn.append(
+                            LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+                        )
+
+            if level != 0:
+                stage.upsample = LTX2AudioUpsample(
+                    block_in, True, causality_axis=self.causality_axis
+                )
+                curr_res *= 2
+
+            self.up.insert(0, stage)
+
+        final_block_channels = block_in
+
+        if self.norm_type == "group":
+            self.norm_out = nn.GroupNorm(
+                num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True
+            )
+        elif self.norm_type == "pixel":
+            self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {self.norm_type}")
+
+        if self.causality_axis is not None:
+            self.conv_out = LTX2AudioCausalConv2d(
+                final_block_channels,
+                output_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_out = nn.Conv2d(
+                final_block_channels,
+                output_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            )
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        _, _, frames, mel_bins = sample.shape
+
+        target_frames = frames * LATENT_DOWNSAMPLE_FACTOR
+
+        if self.causality_axis is not None:
+            target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
+
+        target_channels = self.out_ch
+        target_mel_bins = self.mel_bins if self.mel_bins is not None else mel_bins
+
+        hidden_features = self.conv_in(sample)
+        hidden_features = self.mid.block_1(hidden_features, temb=None)
+        hidden_features = self.mid.attn_1(hidden_features)
+        hidden_features = self.mid.block_2(hidden_features, temb=None)
+
+        for level in reversed(range(self.num_resolutions)):
+            stage = self.up[level]
+            for block_idx, block in enumerate(stage.block):
+                hidden_features = block(hidden_features, temb=None)
+                if stage.attn:
+                    hidden_features = stage.attn[block_idx](hidden_features)
+
+            if level != 0 and hasattr(stage, "upsample"):
+                hidden_features = stage.upsample(hidden_features)
+
+        if self.give_pre_end:
+            return hidden_features
+
+        hidden = self.norm_out(hidden_features)
+        hidden = self.non_linearity(hidden)
+        decoded_output = self.conv_out(hidden)
+        decoded_output = torch.tanh(decoded_output) if self.tanh_out else decoded_output
+
+        _, _, current_time, current_freq = decoded_output.shape
+        target_time = target_frames
+        target_freq = target_mel_bins
+
+        decoded_output = decoded_output[
+            :,
+            :target_channels,
+            : min(current_time, target_time),
+            : min(current_freq, target_freq),
+        ]
+
+        time_padding_needed = target_time - decoded_output.shape[2]
+        freq_padding_needed = target_freq - decoded_output.shape[3]
+
+        if time_padding_needed > 0 or freq_padding_needed > 0:
+            padding = (
+                0,
+                max(freq_padding_needed, 0),
+                0,
+                max(time_padding_needed, 0),
+            )
+            decoded_output = F.pad(decoded_output, padding)
+
+        decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
+
+        return decoded_output
+
+
+class AutoencoderKLLTX2Audio(ParallelTiledVAE):
+    r"""
+    LTX2 audio VAE for encoding and decoding audio latent representations.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    def __init__(
+        self,
+        config: LTXAudioVAEConfig,
+    ) -> None:
+        super().__init__(config=config)
+
+        causality_axis = config.arch_config.causality_axis
+        attn_resolutions = config.arch_config.attn_resolutions
+        base_channels = config.arch_config.base_channels
+        output_channels = config.arch_config.output_channels
+        ch_mult = config.arch_config.ch_mult
+        num_res_blocks = config.arch_config.num_res_blocks
+        in_channels = config.arch_config.in_channels
+        resolution = config.arch_config.resolution
+        latent_channels = config.arch_config.latent_channels
+        norm_type = config.arch_config.norm_type
+        dropout = config.arch_config.dropout
+        mid_block_add_attention = config.arch_config.mid_block_add_attention
+        sample_rate = config.arch_config.sample_rate
+        mel_hop_length = config.arch_config.mel_hop_length
+        is_causal = config.arch_config.is_causal
+        mel_bins = config.arch_config.mel_bins
+        double_z = config.arch_config.double_z
+
+        supported_causality_axes = {"none", "width", "height", "width-compatibility"}
+        if causality_axis not in supported_causality_axes:
+            raise ValueError(
+                f"{causality_axis=} is not valid. Supported values: {supported_causality_axes}"
+            )
+
+        attn_resolution_set = (
+            set(attn_resolutions) if attn_resolutions else attn_resolutions
+        )
+
+        self.encoder = LTX2AudioEncoder(
+            base_channels=base_channels,
+            output_channels=output_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolution_set,
+            in_channels=in_channels,
+            resolution=resolution,
+            latent_channels=latent_channels,
+            norm_type=norm_type,
+            causality_axis=causality_axis,
+            dropout=dropout,
+            mid_block_add_attention=mid_block_add_attention,
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+            double_z=double_z,
+        )
+
+        self.decoder = LTX2AudioDecoder(
+            base_channels=base_channels,
+            output_channels=output_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolution_set,
+            in_channels=in_channels,
+            resolution=resolution,
+            latent_channels=latent_channels,
+            norm_type=norm_type,
+            causality_axis=causality_axis,
+            dropout=dropout,
+            mid_block_add_attention=mid_block_add_attention,
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+        )
+
+        # Per-channel statistics for normalizing and denormalizing the latent representation. This statistics is computed over
+        # the entire dataset and stored in model's checkpoint under AudioVAE state_dict
+        latents_std = torch.ones((base_channels,))
+        latents_mean = torch.zeros((base_channels,))
+        self.register_buffer("latents_mean", latents_mean, persistent=True)
+        self.register_buffer("latents_std", latents_std, persistent=True)
+
+        # TODO: confirm whether the mel compression ratio below is correct
+        self.mel_compression_ratio = LATENT_DOWNSAMPLE_FACTOR
+        self.use_slicing = False
+        # stage containers are not called directly, so hooks attach to called lists
+        self.layer_names = [
+            layer_name
+            for prefix, num_resolutions in (
+                ("encoder.down", self.encoder.num_resolutions),
+                ("decoder.up", self.decoder.num_resolutions),
+            )
+            for level in range(num_resolutions)
+            for layer_name in (
+                f"{prefix}.{level}.block",
+                f"{prefix}.{level}.attn",
+            )
+        ]
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x)
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+
+    def decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z)
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        posterior = self.encode(sample).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        if not return_dict:
+            return (dec.sample,)
+        return dec
+
+
+EntryClass = AutoencoderKLLTX2Audio

python/sglang/multimodal_gen/runtime/models/vaes/parallel/diffusers_spatial.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.parallel_conv import (
+    SpatialParallelConv2d,
+    chunk_height_by_sizes,
+    gather_and_trim_height,
+    gather_variable_height,
+    split_for_parallel_decode,
+)
+
+
+def count_decoder_spatial_upsamples(decoder: nn.Module) -> int:
+    return sum(
+        len(upsamplers)
+        for block in getattr(decoder, "up_blocks", [])
+        if (upsamplers := getattr(block, "upsamplers", None)) is not None
+    )
+
+
+def enable_diffusers_decoder_spatial_parallel(decoder: nn.Module) -> int:
+    _replace_conv2d_modules(decoder)
+    _patch_groupnorm_modules(decoder)
+    _patch_attention_modules(decoder)
+    return count_decoder_spatial_upsamples(decoder)
+
+
+def spatial_parallel_diffusers_decode(
+    decoder: nn.Module, z: torch.Tensor, upsample_count: int
+) -> torch.Tensor:
+    z, expected_height = split_for_parallel_decode(
+        z,
+        upsample_count=upsample_count,
+        world_size=get_decode_parallel_world_size(),
+        rank=get_decode_parallel_rank(),
+    )
+    return gather_and_trim_height(decoder(z), expected_height)
+
+
+def _replace_conv2d_modules(module: nn.Module) -> None:
+    for name, child in list(module.named_children()):
+        if type(child) is nn.Conv2d:
+            setattr(module, name, _make_spatial_conv2d(child))
+        else:
+            _replace_conv2d_modules(child)
+
+
+def _make_spatial_conv2d(conv: nn.Conv2d) -> SpatialParallelConv2d:
+    spatial_conv = SpatialParallelConv2d(
+        in_channels=conv.in_channels,
+        out_channels=conv.out_channels,
+        kernel_size=conv.kernel_size,
+        stride=conv.stride,
+        padding=conv.padding,
+        dilation=conv.dilation,
+        groups=conv.groups,
+        bias=conv.bias is not None,
+        padding_mode=conv.padding_mode,
+    )
+    spatial_conv.weight = conv.weight
+    spatial_conv.bias = conv.bias
+    return spatial_conv
+
+
+def _patch_attention_modules(module: nn.Module) -> None:
+    for child in module.children():
+        if child.__class__.__name__ == "Attention":
+            _patch_attention_forward(child)
+        _patch_attention_modules(child)
+
+
+def _patch_groupnorm_modules(module: nn.Module) -> None:
+    for child in module.children():
+        if type(child) is nn.GroupNorm:
+            _patch_groupnorm_forward(child)
+        else:
+            _patch_groupnorm_modules(child)
+
+
+def _patch_groupnorm_forward(norm: nn.GroupNorm) -> None:
+    original_forward = norm.forward
+
+    def spatial_parallel_forward(hidden_states):
+        if hidden_states.dim() < 4:
+            return original_forward(hidden_states)
+        hidden_states, heights = gather_variable_height(hidden_states)
+        hidden_states = hidden_states.contiguous()
+        hidden_states = original_forward(hidden_states)
+        return chunk_height_by_sizes(hidden_states, heights)
+
+    norm.forward = spatial_parallel_forward
+
+
+def _patch_attention_forward(attn: nn.Module) -> None:
+    original_forward = attn.forward
+
+    def spatial_parallel_forward(hidden_states, *args, **kwargs):
+        if hidden_states.dim() != 4:
+            return original_forward(hidden_states, *args, **kwargs)
+        hidden_states, heights = gather_variable_height(hidden_states)
+        hidden_states = hidden_states.contiguous()
+        hidden_states = original_forward(hidden_states, *args, **kwargs)
+        return chunk_height_by_sizes(hidden_states, heights)
+
+    attn.forward = spatial_parallel_forward

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

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import Pi05PolicyModel
+
+__all__ = ["Pi05PolicyModel"]

python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py

@@ -0,0 +1,759 @@
+import math
+from abc import ABC
+from contextlib import nullcontext
+from typing import Tuple
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+
+LRELU_SLOPE = 0.1
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def _sinc(x: torch.Tensor) -> torch.Tensor:
+    return torch.where(
+        x == 0,
+        torch.tensor(1.0, device=x.device, dtype=x.dtype),
+        torch.sin(math.pi * x) / math.pi / x,
+    )
+
+
+def kaiser_sinc_filter1d(
+    cutoff: float, half_width: float, kernel_size: int
+) -> torch.Tensor:
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+    delta_f = 4 * half_width
+    amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if amplitude > 50.0:
+        beta = 0.1102 * (amplitude - 8.7)
+    elif amplitude >= 21.0:
+        beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+    time = (
+        torch.arange(-half_size, half_size) + 0.5
+        if even
+        else torch.arange(kernel_size) - half_size
+    )
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time)
+        filter_ /= filter_.sum()
+    return filter_.view(1, 1, kernel_size)
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff: float = 0.5,
+        half_width: float = 0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self.register_buffer(
+            "filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, channels, _ = x.shape
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        return F.conv1d(
+            x,
+            self.filter.expand(channels, -1, -1),
+            stride=self.stride,
+            groups=channels,
+        )
+
+
+class UpSample1d(nn.Module):
+    def __init__(
+        self,
+        ratio: int = 2,
+        kernel_size: int | None = None,
+        persistent: bool = True,
+        window_type: str = "kaiser",
+    ):
+        super().__init__()
+        self.ratio = ratio
+        self.stride = ratio
+
+        if window_type == "hann":
+            rolloff = 0.99
+            lowpass_filter_width = 6
+            width = math.ceil(lowpass_filter_width / rolloff)
+            self.kernel_size = 2 * width * ratio + 1
+            self.pad = width
+            self.pad_left = 2 * width * ratio
+            self.pad_right = self.kernel_size - ratio
+            time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff
+            time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width)
+            window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2
+            sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(
+                1, 1, -1
+            )
+        else:
+            self.kernel_size = (
+                int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+            )
+            self.pad = self.kernel_size // ratio - 1
+            self.pad_left = (
+                self.pad * self.stride + (self.kernel_size - self.stride) // 2
+            )
+            self.pad_right = (
+                self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+            )
+            sinc_filter = kaiser_sinc_filter1d(
+                cutoff=0.5 / ratio,
+                half_width=0.6 / ratio,
+                kernel_size=self.kernel_size,
+            )
+
+        self.register_buffer("filter", sinc_filter, persistent=persistent)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, channels, _ = x.shape
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        filt = self.filter.to(dtype=x.dtype, device=x.device).expand(channels, -1, -1)
+        x = self.ratio * F.conv_transpose1d(
+            x, filt, stride=self.stride, groups=channels
+        )
+        return x[..., self.pad_left : -self.pad_right]
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio: int = 2, kernel_size: int | None = None):
+        super().__init__()
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation: nn.Module,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
+        super().__init__()
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.upsample(x)
+        x = self.act(x)
+        return self.downsample(x)
+
+
+class Snake(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        alpha: float = 1.0,
+        alpha_trainable: bool = True,
+        alpha_logscale: bool = True,
+    ):
+        super().__init__()
+        self.alpha_logscale = alpha_logscale
+        self.alpha = nn.Parameter(
+            torch.zeros(in_features)
+            if alpha_logscale
+            else torch.ones(in_features) * alpha
+        )
+        self.alpha.requires_grad = alpha_trainable
+        self.eps = 1e-9
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        return x + (1.0 / (alpha + self.eps)) * torch.sin(x * alpha).pow(2)
+
+
+class SnakeBeta(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        alpha: float = 1.0,
+        alpha_trainable: bool = True,
+        alpha_logscale: bool = True,
+    ):
+        super().__init__()
+        self.alpha_logscale = alpha_logscale
+        self.alpha = nn.Parameter(
+            torch.zeros(in_features)
+            if alpha_logscale
+            else torch.ones(in_features) * alpha
+        )
+        self.alpha.requires_grad = alpha_trainable
+        self.beta = nn.Parameter(
+            torch.zeros(in_features)
+            if alpha_logscale
+            else torch.ones(in_features) * alpha
+        )
+        self.beta.requires_grad = alpha_trainable
+        self.eps = 1e-9
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2)
+
+
+class ResBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        dilations: Tuple[int, ...] = (1, 3, 5),
+        leaky_relu_negative_slope: float = 0.1,
+        padding_mode: str = "same",
+    ):
+        super().__init__()
+        self.dilations = dilations
+        self.negative_slope = leaky_relu_negative_slope
+
+        self.convs1 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=dilation,
+                    padding=padding_mode,
+                )
+                for dilation in dilations
+            ]
+        )
+
+        self.convs2 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=1,
+                    padding=padding_mode,
+                )
+                for _ in range(len(dilations))
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv1, conv2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, negative_slope=self.negative_slope)
+            xt = conv1(xt)
+            xt = F.leaky_relu(xt, negative_slope=self.negative_slope)
+            xt = conv2(xt)
+            x = x + xt
+        return x
+
+
+class AMPBlock1(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple[int, int, int] = (1, 3, 5),
+        activation: str = "snake",
+    ):
+        super().__init__()
+        act_cls = SnakeBeta if activation == "snakebeta" else Snake
+        self.convs1 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[0],
+                    padding=get_padding(kernel_size, dilation[0]),
+                ),
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[1],
+                    padding=get_padding(kernel_size, dilation[1]),
+                ),
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[2],
+                    padding=get_padding(kernel_size, dilation[2]),
+                ),
+            ]
+        )
+        self.convs2 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=1,
+                    padding=get_padding(kernel_size, 1),
+                ),
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=1,
+                    padding=get_padding(kernel_size, 1),
+                ),
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=1,
+                    padding=get_padding(kernel_size, 1),
+                ),
+            ]
+        )
+        self.acts1 = nn.ModuleList(
+            [Activation1d(act_cls(channels)) for _ in range(len(self.convs1))]
+        )
+        self.acts2 = nn.ModuleList(
+            [Activation1d(act_cls(channels)) for _ in range(len(self.convs2))]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv1, conv2, act1, act2 in zip(
+            self.convs1, self.convs2, self.acts1, self.acts2
+        ):
+            xt = act1(x)
+            xt = conv1(xt)
+            xt = act2(xt)
+            xt = conv2(xt)
+            x = x + xt
+        return x
+
+
+class LTX23MelSTFT(nn.Module):
+    class STFTFn(nn.Module):
+        def __init__(self, filter_length: int, hop_length: int, win_length: int):
+            super().__init__()
+            self.hop_length = hop_length
+            self.win_length = win_length
+            n_freqs = filter_length // 2 + 1
+            self.register_buffer(
+                "forward_basis", torch.zeros(n_freqs * 2, 1, filter_length)
+            )
+            self.register_buffer(
+                "inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length)
+            )
+
+        def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            if y.dim() == 2:
+                y = y.unsqueeze(1)
+            left_pad = max(0, self.win_length - self.hop_length)
+            y = F.pad(y, (left_pad, 0))
+            spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0)
+            n_freqs = spec.shape[1] // 2
+            real, imag = spec[:, :n_freqs], spec[:, n_freqs:]
+            magnitude = torch.sqrt(real**2 + imag**2)
+            phase = torch.atan2(imag.float(), real.float()).to(real.dtype)
+            return magnitude, phase
+
+    def __init__(
+        self, filter_length: int, hop_length: int, win_length: int, n_mel_channels: int
+    ):
+        super().__init__()
+        self.stft_fn = self.STFTFn(filter_length, hop_length, win_length)
+        n_freqs = filter_length // 2 + 1
+        self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs))
+
+    def mel_spectrogram(
+        self, y: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        magnitude, phase = self.stft_fn(y)
+        energy = torch.norm(magnitude, dim=1)
+        mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude)
+        log_mel = torch.log(torch.clamp(mel, min=1e-5))
+        return log_mel, magnitude, phase, energy
+
+
+class LTX23VocoderCore(nn.Module):
+    def __init__(  # noqa: PLR0913
+        self,
+        resblock_kernel_sizes: list[int] | None = None,
+        upsample_rates: list[int] | None = None,
+        upsample_kernel_sizes: list[int] | None = None,
+        resblock_dilation_sizes: list[list[int]] | None = None,
+        upsample_initial_channel: int = 1024,
+        resblock: str = "1",
+        output_sampling_rate: int = 24000,
+        activation: str = "snake",
+        use_tanh_at_final: bool = True,
+        apply_final_activation: bool = True,
+        use_bias_at_final: bool = True,
+    ):
+        super().__init__()
+        if resblock_kernel_sizes is None:
+            resblock_kernel_sizes = [3, 7, 11]
+        if upsample_rates is None:
+            upsample_rates = [6, 5, 2, 2, 2]
+        if upsample_kernel_sizes is None:
+            upsample_kernel_sizes = [16, 15, 8, 4, 4]
+        if resblock_dilation_sizes is None:
+            resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+
+        self.output_sampling_rate = output_sampling_rate
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.use_tanh_at_final = use_tanh_at_final
+        self.apply_final_activation = apply_final_activation
+        self.is_amp = resblock == "AMP1"
+
+        self.conv_pre = nn.Conv1d(
+            in_channels=128,
+            out_channels=upsample_initial_channel,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+        )
+        self.ups = nn.ModuleList(
+            nn.ConvTranspose1d(
+                upsample_initial_channel // (2**i),
+                upsample_initial_channel // (2 ** (i + 1)),
+                kernel_size,
+                stride,
+                padding=(kernel_size - stride) // 2,
+            )
+            for i, (stride, kernel_size) in enumerate(
+                zip(upsample_rates, upsample_kernel_sizes, strict=True)
+            )
+        )
+
+        final_channels = upsample_initial_channel // (2 ** len(upsample_rates))
+        self.resblocks = nn.ModuleList()
+        for i in range(len(upsample_rates)):
+            channels = upsample_initial_channel // (2 ** (i + 1))
+            for kernel_size, dilations in zip(
+                resblock_kernel_sizes, resblock_dilation_sizes, strict=True
+            ):
+                if self.is_amp:
+                    self.resblocks.append(
+                        AMPBlock1(
+                            channels,
+                            kernel_size,
+                            tuple(dilations),
+                            activation=activation,
+                        )
+                    )
+                else:
+                    self.resblocks.append(
+                        ResBlock(
+                            channels,
+                            kernel_size=kernel_size,
+                            dilations=tuple(dilations),
+                            leaky_relu_negative_slope=LRELU_SLOPE,
+                            padding_mode=get_padding(kernel_size, 1),
+                        )
+                    )
+
+        self.act_post = (
+            Activation1d(SnakeBeta(final_channels)) if self.is_amp else nn.LeakyReLU()
+        )
+        self.conv_post = nn.Conv1d(
+            in_channels=final_channels,
+            out_channels=2,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+            bias=use_bias_at_final,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(2, 3)
+        if x.dim() == 4:
+            assert x.shape[1] == 2, "Input must have 2 channels for stereo"
+            x = einops.rearrange(x, "b s c t -> b (s c) t")
+
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            if not self.is_amp:
+                x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            start = i * self.num_kernels
+            end = start + self.num_kernels
+            block_outputs = torch.stack(
+                [self.resblocks[idx](x) for idx in range(start, end)],
+                dim=0,
+            )
+            x = block_outputs.mean(dim=0)
+
+        x = self.act_post(x)
+        x = self.conv_post(x)
+        if self.apply_final_activation:
+            x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1)
+        return x
+
+
+class LTX2Vocoder(ABC, nn.Module, LayerwiseOffloadableModuleMixin):
+    r"""
+    LTX 2.0 vocoder for converting generated mel spectrograms back to audio waveforms.
+    """
+
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "upsamplers",
+        "resnets",
+        "vocoder.ups",
+        "vocoder.resblocks",
+        "bwe_generator.ups",
+        "bwe_generator.resblocks",
+    ]
+
+    def __init__(
+        self,
+        config: LTXVocoderConfig,
+    ):
+        super().__init__()
+        self.config = config
+        nested_vocoder_cfg = getattr(config.arch_config, "vocoder", None)
+        if isinstance(nested_vocoder_cfg, dict) and "bwe" in nested_vocoder_cfg:
+            vocoder_cfg = nested_vocoder_cfg.get("vocoder", {})
+            bwe_cfg = nested_vocoder_cfg["bwe"]
+            self.vocoder = LTX23VocoderCore(
+                resblock_kernel_sizes=vocoder_cfg.get("resblock_kernel_sizes"),
+                upsample_rates=vocoder_cfg.get("upsample_rates"),
+                upsample_kernel_sizes=vocoder_cfg.get("upsample_kernel_sizes"),
+                resblock_dilation_sizes=vocoder_cfg.get("resblock_dilation_sizes"),
+                upsample_initial_channel=vocoder_cfg.get(
+                    "upsample_initial_channel", 1024
+                ),
+                resblock=vocoder_cfg.get("resblock", "1"),
+                output_sampling_rate=bwe_cfg["input_sampling_rate"],
+                activation=vocoder_cfg.get("activation", "snake"),
+                use_tanh_at_final=vocoder_cfg.get("use_tanh_at_final", True),
+                apply_final_activation=vocoder_cfg.get("apply_final_activation", True),
+                use_bias_at_final=vocoder_cfg.get("use_bias_at_final", True),
+            )
+            self.bwe_generator = LTX23VocoderCore(
+                resblock_kernel_sizes=bwe_cfg.get("resblock_kernel_sizes"),
+                upsample_rates=bwe_cfg.get("upsample_rates"),
+                upsample_kernel_sizes=bwe_cfg.get("upsample_kernel_sizes"),
+                resblock_dilation_sizes=bwe_cfg.get("resblock_dilation_sizes"),
+                upsample_initial_channel=bwe_cfg.get("upsample_initial_channel", 1024),
+                resblock=bwe_cfg.get("resblock", "1"),
+                output_sampling_rate=bwe_cfg["output_sampling_rate"],
+                activation=bwe_cfg.get("activation", "snake"),
+                use_tanh_at_final=bwe_cfg.get("use_tanh_at_final", True),
+                apply_final_activation=bwe_cfg.get("apply_final_activation", True),
+                use_bias_at_final=bwe_cfg.get("use_bias_at_final", True),
+            )
+            self.mel_stft = LTX23MelSTFT(
+                filter_length=bwe_cfg["n_fft"],
+                hop_length=bwe_cfg["hop_length"],
+                win_length=bwe_cfg.get("win_size", bwe_cfg["n_fft"]),
+                n_mel_channels=bwe_cfg["num_mels"],
+            )
+            self.input_sampling_rate = bwe_cfg["input_sampling_rate"]
+            self.output_sampling_rate = bwe_cfg["output_sampling_rate"]
+            self.hop_length = bwe_cfg["hop_length"]
+            with torch.device("cpu"):
+                self.resampler = UpSample1d(
+                    ratio=self.output_sampling_rate // self.input_sampling_rate,
+                    persistent=False,
+                    window_type="hann",
+                )
+            self.sample_rate = self.output_sampling_rate
+            return
+
+        self.sample_rate = (
+            getattr(config.arch_config, "sample_rate", None)
+            or getattr(config.arch_config, "sampling_rate", None)
+            or getattr(config.arch_config, "audio_sample_rate", None)
+            or getattr(config.arch_config, "output_sampling_rate", None)
+        )
+
+        in_channels = config.arch_config.in_channels
+        hidden_channels = config.arch_config.hidden_channels
+        out_channels = config.arch_config.out_channels
+        upsample_kernel_sizes = config.arch_config.upsample_kernel_sizes
+        upsample_factors = config.arch_config.upsample_factors
+        resnet_kernel_sizes = config.arch_config.resnet_kernel_sizes
+        resnet_dilations = config.arch_config.resnet_dilations
+        leaky_relu_negative_slope = config.arch_config.leaky_relu_negative_slope
+
+        self.num_upsample_layers = len(upsample_kernel_sizes)
+        self.resnets_per_upsample = len(resnet_kernel_sizes)
+        self.out_channels = out_channels
+        self.total_upsample_factor = math.prod(upsample_factors)
+        self.negative_slope = leaky_relu_negative_slope
+
+        if self.num_upsample_layers != len(upsample_factors):
+            raise ValueError(
+                f"`upsample_kernel_sizes` and `upsample_factors` should be lists of the same length but are length"
+                f" {self.num_upsample_layers} and {len(upsample_factors)}, respectively."
+            )
+
+        if self.resnets_per_upsample != len(resnet_dilations):
+            raise ValueError(
+                f"`resnet_kernel_sizes` and `resnet_dilations` should be lists of the same length but are length"
+                f" {len(self.resnets_per_upsample)} and {len(resnet_dilations)}, respectively."
+            )
+
+        self.conv_in = nn.Conv1d(
+            in_channels, hidden_channels, kernel_size=7, stride=1, padding=3
+        )
+
+        self.upsamplers = nn.ModuleList()
+        self.resnets = nn.ModuleList()
+        input_channels = hidden_channels
+        for i, (stride, kernel_size) in enumerate(
+            zip(upsample_factors, upsample_kernel_sizes)
+        ):
+            output_channels = input_channels // 2
+            self.upsamplers.append(
+                nn.ConvTranspose1d(
+                    input_channels,  # hidden_channels // (2 ** i)
+                    output_channels,  # hidden_channels // (2 ** (i + 1))
+                    kernel_size,
+                    stride=stride,
+                    padding=(kernel_size - stride) // 2,
+                )
+            )
+
+            for kernel_size, dilations in zip(resnet_kernel_sizes, resnet_dilations):
+                self.resnets.append(
+                    ResBlock(
+                        output_channels,
+                        kernel_size,
+                        dilations=dilations,
+                        leaky_relu_negative_slope=leaky_relu_negative_slope,
+                    )
+                )
+            input_channels = output_channels
+
+        self.conv_out = nn.Conv1d(output_channels, out_channels, 7, stride=1, padding=3)
+
+    def _compute_ltx23_mel(self, audio: torch.Tensor) -> torch.Tensor:
+        batch, channels, _ = audio.shape
+        flat = audio.reshape(batch * channels, -1)
+        mel, _, _, _ = self.mel_stft.mel_spectrogram(flat)
+        return mel.reshape(batch, channels, mel.shape[1], mel.shape[2])
+
+    def forward(
+        self, hidden_states: torch.Tensor, time_last: bool = False
+    ) -> torch.Tensor:
+        r"""
+        Forward pass of the vocoder.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                Input Mel spectrogram tensor of shape `(batch_size, num_channels, time, num_mel_bins)` if `time_last`
+                is `False` (the default) or shape `(batch_size, num_channels, num_mel_bins, time)` if `time_last` is
+                `True`.
+            time_last (`bool`, *optional*, defaults to `False`):
+                Whether the last dimension of the input is the time/frame dimension or the Mel bins dimension.
+
+        Returns:
+            `torch.Tensor`:
+                Audio waveform tensor of shape (batch_size, out_channels, audio_length)
+        """
+        if hasattr(self, "bwe_generator"):
+            input_dtype = hidden_states.dtype
+            autocast_ctx = (
+                torch.autocast(
+                    device_type=hidden_states.device.type, dtype=torch.float32
+                )
+                if hidden_states.device.type != "cpu"
+                else nullcontext()
+            )
+            with autocast_ctx:
+                waveform = self.vocoder(hidden_states.float())
+                length_low_rate = waveform.shape[-1]
+                output_length = (
+                    length_low_rate
+                    * self.output_sampling_rate
+                    // self.input_sampling_rate
+                )
+                remainder = length_low_rate % self.hop_length
+                if remainder != 0:
+                    waveform = F.pad(waveform, (0, self.hop_length - remainder))
+                mel = self._compute_ltx23_mel(waveform)
+                residual = self.bwe_generator(mel.transpose(2, 3))
+                skip = self.resampler(waveform)
+                assert residual.shape == skip.shape
+                waveform = torch.clamp(residual + skip, -1, 1)[..., :output_length]
+            return waveform.to(input_dtype)
+
+        # Ensure that the time/frame dimension is last
+        if not time_last:
+            hidden_states = hidden_states.transpose(2, 3)
+        # Combine channels and frequency (mel bins) dimensions
+        hidden_states = hidden_states.flatten(1, 2)
+
+        hidden_states = self.conv_in(hidden_states)
+
+        for i in range(self.num_upsample_layers):
+            hidden_states = F.leaky_relu(
+                hidden_states, negative_slope=self.negative_slope
+            )
+            hidden_states = self.upsamplers[i](hidden_states)
+
+            # Run all resnets in parallel on hidden_states
+            start = i * self.resnets_per_upsample
+            end = (i + 1) * self.resnets_per_upsample
+            resnet_outputs = torch.stack(
+                [self.resnets[j](hidden_states) for j in range(start, end)], dim=0
+            )
+
+            hidden_states = torch.mean(resnet_outputs, dim=0)
+
+        # NOTE: unlike the first leaky ReLU, this leaky ReLU is set to use the default F.leaky_relu negative slope of
+        # 0.01 (whereas the others usually use a slope of 0.1). Not sure if this is intended
+        hidden_states = F.leaky_relu(hidden_states, negative_slope=0.01)
+        hidden_states = self.conv_out(hidden_states)
+        hidden_states = torch.tanh(hidden_states)
+
+        return hidden_states
+
+
+EntryClass = LTX2Vocoder