chore: import upstream snapshot with attribution

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