chore: import upstream snapshot with attribution

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]