sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/dits/zimage.py

import math
from typing import Any, List, Optional, Tuple

import torch
import torch.nn as nn

from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
from sglang.multimodal_gen.runtime.distributed import (
    get_sp_parallel_rank,
    get_sp_world_size,
    get_tp_world_size,
    sequence_model_parallel_all_gather,
)
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
    get_ring_parallel_world_size,
)
from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
from sglang.multimodal_gen.runtime.layers.attention import (
    UlyssesAttention,
    USPAttention,
    build_varlen_mask_meta_from_lengths,
    build_varlen_mask_meta_from_ranges,
)
from sglang.multimodal_gen.runtime.layers.layernorm import (
    apply_qk_norm_with_optional_rope,
)
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.configs.nunchaku_config import (
    NunchakuConfig,
    is_nunchaku_available,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    _apply_rotary_emb,
    apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

try:
    from nunchaku.models.attention import NunchakuFeedForward  # type: ignore[import]
except Exception:
    NunchakuFeedForward = None

logger = init_logger(__name__)
_is_cuda = current_platform.is_cuda()

ADALN_EMBED_DIM = 256
SEQ_MULTI_OF = 32


class ZImageRMSNorm(nn.Module):
    """RMSNorm that preserves Z-Image's native bf16 behavior.

    Z-Image does not upcast hidden states to fp32 for RMSNorm.
    """

    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.variance_epsilon = eps
        self.hidden_size = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_dtype = x.dtype
        output = x * torch.rsqrt(
            x.pow(2).mean(dim=-1, keepdim=True) + self.variance_epsilon
        )
        output = output * self.weight.to(device=x.device, dtype=x.dtype)
        return output.to(orig_dtype)


def zimage_rmsnorm_tanh_mul_add(
    x: torch.Tensor,
    gate: torch.Tensor,
    residual: torch.Tensor,
    norm: ZImageRMSNorm,
    enable_fused: bool = True,
) -> torch.Tensor:
    if enable_fused:
        from sglang.jit_kernel.diffusion.triton.zimage_native_norm import (
            zimage_rmsnorm_tanh_residual,
        )

        y = zimage_rmsnorm_tanh_residual(
            x,
            gate,
            residual,
            norm.weight.data.to(device=x.device, dtype=x.dtype).contiguous(),
            norm.variance_epsilon,
        )
        if y is not None:
            return y
    return residual + torch.tanh(gate) * norm(x)


def zimage_rmsnorm_scale(
    x: torch.Tensor,
    scale: torch.Tensor,
    norm: ZImageRMSNorm,
    enable_fused: bool = True,
) -> torch.Tensor:
    if enable_fused:
        from sglang.jit_kernel.diffusion.triton.zimage_native_norm import (
            zimage_rmsnorm_scale as fused_zimage_rmsnorm_scale,
        )

        y = fused_zimage_rmsnorm_scale(
            x,
            norm.weight.data.to(device=x.device, dtype=x.dtype).contiguous(),
            scale,
            norm.variance_epsilon,
        )
        if y is not None:
            return y
    return norm(x) * scale


class SelectFirstElement(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x[0]


class TimestepEmbedder(nn.Module):
    def __init__(self, out_size, mid_size=None, frequency_embedding_size=256):
        super().__init__()
        if mid_size is None:
            mid_size = out_size

        self.mlp = nn.ModuleList(
            [
                ColumnParallelLinear(
                    frequency_embedding_size, mid_size, bias=True, gather_output=False
                ),
                nn.SiLU(),
                RowParallelLinear(
                    mid_size, out_size, bias=True, input_is_parallel=True
                ),
            ]
        )

        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        with torch.amp.autocast(current_platform.device_type, enabled=False):
            half = dim // 2
            freqs = torch.exp(
                -math.log(max_period)
                * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
                / half
            )
            args = t[:, None].float() * freqs[None]
            embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
            if dim % 2:
                embedding = torch.cat(
                    [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
                )
            return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(
            self.mlp[0].weight.dtype
        )
        t_emb, _ = self.mlp[0](t_freq)
        t_emb = self.mlp[1](t_emb)
        t_emb, _ = self.mlp[2](t_emb)
        return t_emb


class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        # Use MergedColumnParallelLinear for gate and up projection (fused)
        self.w13 = MergedColumnParallelLinear(
            dim,
            [hidden_dim, hidden_dim],
            bias=False,
            gather_output=False,
            quant_config=quant_config,
            prefix=f"{prefix}.w13",
        )
        self.w2 = RowParallelLinear(
            hidden_dim,
            dim,
            bias=False,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=f"{prefix}.w2",
        )
        self.act = SiluAndMul()

    def forward(self, x):
        x13, _ = self.w13(x)
        x = self.act(x13)
        out, _ = self.w2(x)
        return out


class ZImageAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        num_kv_heads: int,
        qk_norm: bool = True,
        eps: float = 1e-6,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.dim = dim
        self.head_dim = dim // num_heads
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.qk_norm = qk_norm
        self.enable_zimage_qk_fusion = quant_config is None

        tp_size = get_tp_world_size()
        assert (
            num_heads % tp_size == 0
        ), f"num_heads {num_heads} must be divisible by tp world size {tp_size}"
        assert (
            num_kv_heads % tp_size == 0
        ), f"num_kv_heads {num_kv_heads} must be divisible by tp world size {tp_size}"
        self.local_num_heads = num_heads // tp_size
        self.local_num_kv_heads = num_kv_heads // tp_size

        kv_dim = self.head_dim * num_kv_heads
        self.use_fused_qkv = True

        if self.use_fused_qkv:
            self.to_qkv = MergedColumnParallelLinear(
                dim,
                [dim, kv_dim, kv_dim],
                bias=False,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_qkv",
            )
        else:
            self.to_q = ColumnParallelLinear(
                dim,
                dim,
                bias=False,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_q",
            )
            self.to_k = ColumnParallelLinear(
                dim,
                kv_dim,
                bias=False,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_k",
            )
            self.to_v = ColumnParallelLinear(
                dim,
                kv_dim,
                bias=False,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_v",
            )

        if self.qk_norm:
            self.norm_q = ZImageRMSNorm(self.head_dim, eps=eps)
            self.norm_k = ZImageRMSNorm(self.head_dim, eps=eps)
        else:
            self.norm_q = None
            self.norm_k = None

        self.to_out = nn.ModuleList(
            [
                RowParallelLinear(
                    dim,
                    dim,
                    bias=False,
                    input_is_parallel=True,
                    quant_config=quant_config,
                    prefix=f"{prefix}.to_out.0",
                )
            ]
        )

        self.attn = USPAttention(
            num_heads=self.local_num_heads,
            head_size=self.head_dim,
            num_kv_heads=self.local_num_kv_heads,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
        )
        self.ulysses_attn = UlyssesAttention(
            num_heads=self.local_num_heads,
            head_size=self.head_dim,
            num_kv_heads=self.local_num_kv_heads,
            softmax_scale=None,
            causal=False,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        rope_cos_sin_cache: Optional[torch.Tensor] = None,
        rope_positions: Optional[torch.Tensor] = None,
        attn_mask: Optional[torch.Tensor] = None,
        attn_mask_meta: Optional[dict] = None,
        num_replicated_prefix: int = 0,
        num_replicated_suffix: int = 0,
        skip_sequence_parallel_override: bool = False,
    ):
        if self.use_fused_qkv:
            qkv, _ = self.to_qkv(hidden_states)
            q, k, v = qkv.split(
                [
                    self.local_num_heads * self.head_dim,
                    self.local_num_kv_heads * self.head_dim,
                    self.local_num_kv_heads * self.head_dim,
                ],
                dim=-1,
            )
            q = q.contiguous()
            k = k.contiguous()
            v = v.contiguous()
        else:
            q, _ = self.to_q(hidden_states)
            k, _ = self.to_k(hidden_states)
            v, _ = self.to_v(hidden_states)
        q = q.view(*q.shape[:-1], self.local_num_heads, self.head_dim)
        k = k.view(*k.shape[:-1], self.local_num_kv_heads, self.head_dim)
        v = v.view(*v.shape[:-1], self.local_num_kv_heads, self.head_dim)

        if rope_cos_sin_cache is not None:
            if self.qk_norm:
                q, k = apply_qk_norm_with_optional_rope(
                    q=q,
                    k=k,
                    q_norm=self.norm_q,
                    k_norm=self.norm_k,
                    head_dim=self.head_dim,
                    cos_sin_cache=rope_cos_sin_cache,
                    is_neox=False,
                    positions=rope_positions,
                    allow_inplace=False,
                )
            else:
                q, k = apply_flashinfer_rope_qk_inplace(
                    q,
                    k,
                    rope_cos_sin_cache,
                    is_neox=False,
                    positions=rope_positions,
                )
        elif freqs_cis is not None:
            cos, sin = freqs_cis
            if cos.dim() == 3:
                batch_size, seq_len = q.shape[:2]
                cos_sin_cache = torch.cat(
                    [
                        cos.to(dtype=torch.float32).contiguous(),
                        sin.to(dtype=torch.float32).contiguous(),
                    ],
                    dim=-1,
                ).reshape(batch_size * seq_len, -1)
                positions = torch.arange(
                    batch_size * seq_len, device=q.device, dtype=torch.long
                )
                if self.qk_norm:
                    q, k = apply_qk_norm_with_optional_rope(
                        q=q,
                        k=k,
                        q_norm=self.norm_q,
                        k_norm=self.norm_k,
                        head_dim=self.head_dim,
                        cos_sin_cache=cos_sin_cache,
                        is_neox=False,
                        positions=positions,
                        allow_inplace=self.enable_zimage_qk_fusion,
                    )
                else:
                    q, k = apply_flashinfer_rope_qk_inplace(
                        q, k, cos_sin_cache, is_neox=False, positions=positions
                    )
            elif _is_cuda and q.shape == k.shape:
                cos_sin_cache = torch.cat(
                    [
                        cos.to(dtype=torch.float32).contiguous(),
                        sin.to(dtype=torch.float32).contiguous(),
                    ],
                    dim=-1,
                )
                if self.qk_norm:
                    q, k = apply_qk_norm_with_optional_rope(
                        q=q,
                        k=k,
                        q_norm=self.norm_q,
                        k_norm=self.norm_k,
                        head_dim=self.head_dim,
                        cos_sin_cache=cos_sin_cache,
                        is_neox=False,
                        allow_inplace=self.enable_zimage_qk_fusion,
                    )
                else:
                    q, k = apply_flashinfer_rope_qk_inplace(
                        q, k, cos_sin_cache, is_neox=False
                    )
            else:
                if self.qk_norm:
                    q, k = apply_qk_norm_with_optional_rope(
                        q=q,
                        k=k,
                        q_norm=self.norm_q,
                        k_norm=self.norm_k,
                        head_dim=self.head_dim,
                        allow_inplace=self.enable_zimage_qk_fusion,
                    )
                q = _apply_rotary_emb(q, cos, sin, is_neox_style=False)
                k = _apply_rotary_emb(k, cos, sin, is_neox_style=False)
        elif self.qk_norm:
            q, k = apply_qk_norm_with_optional_rope(
                q=q,
                k=k,
                q_norm=self.norm_q,
                k_norm=self.norm_k,
                head_dim=self.head_dim,
                allow_inplace=self.enable_zimage_qk_fusion,
            )

        if (
            num_replicated_suffix > 0
            and get_sp_world_size() > 1
            and get_ring_parallel_world_size() == 1
        ):
            # the cap (last num_replicated_suffix tokens), as condition, should be replicated
            q_shard, q_rep = (
                q[:, :-num_replicated_suffix],
                q[:, -num_replicated_suffix:],
            )
            k_shard, k_rep = (
                k[:, :-num_replicated_suffix],
                k[:, -num_replicated_suffix:],
            )
            v_shard, v_rep = (
                v[:, :-num_replicated_suffix],
                v[:, -num_replicated_suffix:],
            )
            hidden_states, hidden_states_rep = self.ulysses_attn(
                q_shard,
                k_shard,
                v_shard,
                replicated_q=q_rep,
                replicated_k=k_rep,
                replicated_v=v_rep,
            )
            assert hidden_states_rep is not None
            hidden_states = torch.cat([hidden_states, hidden_states_rep], dim=1)
        else:
            hidden_states = self.attn(
                q,
                k,
                v,
                attn_mask=attn_mask,
                attn_mask_meta=attn_mask_meta,
                num_replicated_prefix=num_replicated_prefix,
                num_replicated_suffix=num_replicated_suffix,
                skip_sequence_parallel_override=skip_sequence_parallel_override,
            )
        hidden_states = hidden_states.flatten(2)

        hidden_states, _ = self.to_out[0](hidden_states)

        return hidden_states


class ZImageTransformerBlock(nn.Module):
    def __init__(
        self,
        layer_id: int,
        dim: int,
        n_heads: int,
        n_kv_heads: int,
        norm_eps: float,
        qk_norm: bool,
        modulation=True,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.dim = dim
        self.head_dim = dim // n_heads
        self.layer_id = layer_id
        self.modulation = modulation
        self.enable_zimage_native_norm_fusion = quant_config is None

        self.attention = ZImageAttention(
            dim=dim,
            num_heads=n_heads,
            num_kv_heads=n_kv_heads,
            qk_norm=qk_norm,
            eps=1e-5,
            quant_config=quant_config,
            prefix=f"{prefix}.attention",
        )
        if not modulation:
            # Context refiner runs on fully replicated caption tokens only.
            # Bypass Ulysses here to preserve the single-GPU attention semantics.
            self.attention.attn.skip_sequence_parallel = True

        hidden_dim = int(dim / 3 * 8)
        nunchaku_enabled = (
            isinstance(quant_config, NunchakuConfig) and is_nunchaku_available()
        )
        if nunchaku_enabled:
            import diffusers

            ff = diffusers.models.attention.FeedForward(
                dim=dim,
                dim_out=dim,
                activation_fn="swiglu",
                inner_dim=hidden_dim,
                bias=False,
            )
            nunchaku_kwargs = {
                "precision": quant_config.precision,
                "rank": quant_config.rank,
                "act_unsigned": quant_config.act_unsigned,
            }
            self.feed_forward = NunchakuFeedForward(ff, **nunchaku_kwargs)
            # NunchakuFeedForward overrides net[2].act_unsigned=True for int4 (GELU-specific
            # optimization for non-negative activations). Z-Image uses SwiGLU whose output
            # can be negative, so we must restore the original act_unsigned value.
            if hasattr(self.feed_forward, "net") and len(self.feed_forward.net) > 2:
                self.feed_forward.net[2].act_unsigned = quant_config.act_unsigned
        else:
            self.feed_forward = FeedForward(
                dim=dim,
                hidden_dim=hidden_dim,
                quant_config=quant_config,
                prefix=f"{prefix}.feed_forward",
            )

        self.attention_norm1 = ZImageRMSNorm(dim, eps=norm_eps)
        self.ffn_norm1 = ZImageRMSNorm(dim, eps=norm_eps)

        self.attention_norm2 = ZImageRMSNorm(dim, eps=norm_eps)
        self.ffn_norm2 = ZImageRMSNorm(dim, eps=norm_eps)

        if modulation:
            self.adaLN_modulation = nn.Sequential(
                ReplicatedLinear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True)
            )

    def forward(
        self,
        x: torch.Tensor,
        freqs_cis: Tuple[torch.Tensor, torch.Tensor],
        adaln_input: Optional[torch.Tensor] = None,
        rope_cos_sin_cache: Optional[torch.Tensor] = None,
        rope_positions: Optional[torch.Tensor] = None,
        attn_mask: Optional[torch.Tensor] = None,
        attn_mask_meta: Optional[dict] = None,
        num_replicated_prefix: int = 0,
        num_replicated_suffix: int = 0,
        skip_sequence_parallel_override: bool = False,
    ):
        if self.modulation:
            assert adaln_input is not None
            scale_msa_gate, _ = self.adaLN_modulation(adaln_input)
            scale_msa, gate_msa, scale_mlp, gate_mlp = scale_msa_gate.unsqueeze(
                1
            ).chunk(4, dim=2)
            scale_msa = 1.0 + scale_msa

            # Attention block
            attn_out = self.attention(
                zimage_rmsnorm_scale(
                    x,
                    scale_msa,
                    self.attention_norm1,
                    self.enable_zimage_native_norm_fusion,
                ),
                freqs_cis=freqs_cis,
                rope_cos_sin_cache=rope_cos_sin_cache,
                rope_positions=rope_positions,
                attn_mask=attn_mask,
                attn_mask_meta=attn_mask_meta,
                num_replicated_prefix=num_replicated_prefix,
                num_replicated_suffix=num_replicated_suffix,
                skip_sequence_parallel_override=skip_sequence_parallel_override,
            )
            x = zimage_rmsnorm_tanh_mul_add(
                attn_out,
                gate_msa,
                x,
                self.attention_norm2,
                self.enable_zimage_native_norm_fusion,
            )
            ffn_in = zimage_rmsnorm_scale(
                x,
                1.0 + scale_mlp,
                self.ffn_norm1,
                self.enable_zimage_native_norm_fusion,
            )

            # FFN block
            ffn_out = self.feed_forward(ffn_in)
            x = zimage_rmsnorm_tanh_mul_add(
                ffn_out,
                gate_mlp,
                x,
                self.ffn_norm2,
                self.enable_zimage_native_norm_fusion,
            )
        else:
            # Attention block
            attn_input = self.attention_norm1(x)
            attn_out = self.attention(
                attn_input,
                freqs_cis=freqs_cis,
                rope_cos_sin_cache=rope_cos_sin_cache,
                rope_positions=rope_positions,
                attn_mask=attn_mask,
                attn_mask_meta=attn_mask_meta,
                num_replicated_prefix=num_replicated_prefix,
                num_replicated_suffix=num_replicated_suffix,
                skip_sequence_parallel_override=skip_sequence_parallel_override,
            )
            x = x + self.attention_norm2(attn_out)

            # FFN block
            ffn_input = self.ffn_norm1(x)
            ffn_out = self.feed_forward(
                ffn_input,
            )
            x = x + self.ffn_norm2(ffn_out)

        return x


class FinalLayer(nn.Module):
    def __init__(self, hidden_size, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = ColumnParallelLinear(
            hidden_size, out_channels, bias=True, gather_output=True
        )

        self.act = nn.SiLU()
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            ReplicatedLinear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True),
        )

    def forward(self, x, c):
        scale, _ = self.adaLN_modulation(c)
        scale = 1.0 + scale
        x = self.norm_final(x) * scale.unsqueeze(1)
        x, _ = self.linear(x)
        return x


class RopeEmbedder:
    def __init__(
        self,
        theta: float = 256.0,
        axes_dims: List[int] = (16, 56, 56),
        axes_lens: List[int] = (64, 128, 128),
    ):
        self.theta = theta
        self.axes_dims = axes_dims
        self.axes_lens = axes_lens
        assert len(axes_dims) == len(
            axes_lens
        ), "axes_dims and axes_lens must have the same length"

        self.cos_cached = None
        self.sin_cached = None

    @staticmethod
    def precompute_freqs(dim: List[int], end: List[int], theta: float = 256.0):
        with torch.device("cpu"):
            cos_list = []
            sin_list = []
            for i, (d, e) in enumerate(zip(dim, end)):
                freqs = 1.0 / (
                    theta
                    ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)
                )
                timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
                freqs = torch.outer(timestep, freqs).float()

                cos_list.append(torch.cos(freqs))
                sin_list.append(torch.sin(freqs))

            return cos_list, sin_list

    def __call__(self, ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            ids: [batch, len(axes_dims)] or [seq_len, len(axes_dims)]
        Returns:
            cos: [batch/seq, head_dim // 2]
            sin: [batch/seq, head_dim // 2]
        """
        assert ids.ndim == 2
        assert ids.shape[-1] == len(self.axes_dims)
        device = ids.device

        if self.cos_cached is None:
            self.cos_cached, self.sin_cached = self.precompute_freqs(
                self.axes_dims, self.axes_lens, theta=self.theta
            )
            self.cos_cached = [c.to(device) for c in self.cos_cached]
            self.sin_cached = [s.to(device) for s in self.sin_cached]
        else:
            if self.cos_cached[0].device != device:
                self.cos_cached = [c.to(device) for c in self.cos_cached]
                self.sin_cached = [s.to(device) for s in self.sin_cached]

        cos_out = []
        sin_out = []
        for i in range(len(self.axes_dims)):
            index = ids[:, i]
            cos_out.append(self.cos_cached[i][index])
            sin_out.append(self.sin_cached[i][index])

        return torch.cat(cos_out, dim=-1), torch.cat(sin_out, dim=-1)


class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
    _supports_gradient_checkpointing = True
    _no_split_modules = ["ZImageTransformerBlock"]
    _fsdp_shard_conditions = ZImageDitConfig().arch_config._fsdp_shard_conditions
    param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping

    param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping
    reverse_param_names_mapping = (
        ZImageDitConfig().arch_config.reverse_param_names_mapping
    )

    # Maps fused runtime layer names to their checkpoint shard names.
    # Used by is_layer_skipped() to correctly handle --quantization-ignored-layers
    # Only list fusions that are unconditional. Conditional fusions (e.g. to_qkv for
    # Nunchaku) are handled by their own quant path.
    packed_modules_mapping = {
        "w13": ["w1", "w3"],
    }

    @classmethod
    def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]:
        return {
            "skip": [
                "norm",
                "embed",
                "rotary",
                "pos_embed",
            ],
            "svdq_w4a4": [
                "attention.to_qkv",
                "attention.to_out",
                "img_mlp",
                "txt_mlp",
            ],
            "awq_w4a16": [
                "img_mod",
                "txt_mod",
            ],
        }

    def __init__(
        self,
        config: ZImageDitConfig,
        hf_config: dict[str, Any],
        quant_config: Optional[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.all_patch_size = arch_config.all_patch_size
        self.all_f_patch_size = arch_config.all_f_patch_size
        self.dim = arch_config.dim
        self.n_heads = arch_config.num_attention_heads

        self.rope_theta = arch_config.rope_theta
        self.t_scale = arch_config.t_scale
        self.gradient_checkpointing = False

        assert len(self.all_patch_size) == len(self.all_f_patch_size)

        all_x_embedder = {}
        all_final_layer = {}
        for patch_idx, (patch_size, f_patch_size) in enumerate(
            zip(self.all_patch_size, self.all_f_patch_size)
        ):
            x_embedder = ColumnParallelLinear(
                f_patch_size * patch_size * patch_size * self.in_channels,
                self.dim,
                bias=True,
                gather_output=True,
            )
            all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder

            final_layer = FinalLayer(
                self.dim, patch_size * patch_size * f_patch_size * self.out_channels
            )
            all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer

        self.all_x_embedder = nn.ModuleDict(all_x_embedder)
        self.all_final_layer = nn.ModuleDict(all_final_layer)

        self.noise_refiner = nn.ModuleList(
            [
                ZImageTransformerBlock(
                    1000 + layer_id,
                    self.dim,
                    self.n_heads,
                    arch_config.n_kv_heads,
                    arch_config.norm_eps,
                    arch_config.qk_norm,
                    modulation=True,
                    quant_config=quant_config,
                    prefix=f"noise_refiner.{layer_id}",
                )
                for layer_id in range(arch_config.n_refiner_layers)
            ]
        )
        self.context_refiner = nn.ModuleList(
            [
                ZImageTransformerBlock(
                    layer_id,
                    self.dim,
                    self.n_heads,
                    arch_config.n_kv_heads,
                    arch_config.norm_eps,
                    arch_config.qk_norm,
                    modulation=False,
                    quant_config=quant_config,
                    prefix=f"context_refiner.{layer_id}",
                )
                for layer_id in range(arch_config.n_refiner_layers)
            ]
        )
        self.t_embedder = TimestepEmbedder(
            min(self.dim, ADALN_EMBED_DIM), mid_size=1024
        )

        self.cap_embedder = nn.Sequential(
            ZImageRMSNorm(arch_config.cap_feat_dim, eps=arch_config.norm_eps),
            ReplicatedLinear(arch_config.cap_feat_dim, self.dim, bias=True),
        )

        self.x_pad_token = nn.Parameter(torch.empty((1, self.dim)))
        self.cap_pad_token = nn.Parameter(torch.empty((1, self.dim)))

        self.layers = nn.ModuleList(
            [
                ZImageTransformerBlock(
                    layer_id,
                    self.dim,
                    self.n_heads,
                    arch_config.n_kv_heads,
                    arch_config.norm_eps,
                    arch_config.qk_norm,
                    quant_config=quant_config,
                    prefix=f"layers.{layer_id}",
                )
                for layer_id in range(arch_config.num_layers)
            ]
        )
        head_dim = self.dim // self.n_heads
        assert head_dim == sum(arch_config.axes_dims)
        self.axes_dims = arch_config.axes_dims
        self.axes_lens = arch_config.axes_lens

        self.rotary_emb = RopeEmbedder(
            theta=self.rope_theta, axes_dims=self.axes_dims, axes_lens=self.axes_lens
        )
        self.layer_names = ["layers"]

    def unpatchify(
        self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size
    ) -> List[torch.Tensor]:
        pH = pW = patch_size
        pF = f_patch_size
        bsz = len(x)
        assert len(size) == bsz
        for i in range(bsz):
            F, H, W = size[i]
            ori_len = (F // pF) * (H // pH) * (W // pW)
            # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)"
            x[i] = (
                x[i][:ori_len]
                .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels)
                .permute(6, 0, 3, 1, 4, 2, 5)
                .reshape(self.out_channels, F, H, W)
            )
        return x

    @staticmethod
    def create_coordinate_grid(size, start=None, device=None):
        if start is None:
            start = (0 for _ in size)

        axes = [
            torch.arange(x0, x0 + span, dtype=torch.int32, device=device)
            for x0, span in zip(start, size)
        ]
        grids = torch.meshgrid(axes, indexing="ij")
        return torch.stack(grids, dim=-1)

    @staticmethod
    def _ceil_to_multiple(value: int, multiple: int) -> int:
        if multiple <= 0:
            return value
        return int(math.ceil(value / multiple) * multiple)

    def patchify_and_embed(
        self,
        all_image: List[torch.Tensor],
        all_cap_feats: List[torch.Tensor],
        patch_size: int,
        f_patch_size: int,
        image_seq_len_target: int | None = None,
        caption_valid_lens: torch.Tensor | None = None,
        caption_valid_mask: torch.Tensor | None = None,
    ):
        """Patchify images and pad image/caption tokens to batch targets.

        Each image is [C, F, H, W] and has one [L, D] caption. Returned tensors
        are stacked as [B, S, D], while valid lengths keep track of real tokens
        before learned pad tokens are restored. `image_seq_len_target`, when
        set, is the SP-local padded image-token target.
        """
        if len(all_image) != len(all_cap_feats):
            raise ValueError(
                f"Z-Image expects one caption embedding per image, got {len(all_image)} images and {len(all_cap_feats)} captions"
            )
        if not all_image:
            raise ValueError("Z-Image batch must contain at least one image latent")
        if caption_valid_mask is not None and caption_valid_mask.shape[0] != len(
            all_cap_feats
        ):
            raise ValueError("caption_valid_mask must have one row per Z-Image caption")

        pH = pW = patch_size
        pF = f_patch_size
        all_image_out = []
        all_image_size = []
        all_cap_feats_out = []
        all_image_valid_lens = []
        all_cap_valid_lens = []
        all_cap_valid_masks = []
        all_image_attn_lens = []
        all_cap_attn_lens = []
        image_records = []

        cap_seq_len_target = max(
            self._ceil_to_multiple(cap_feat.size(0), SEQ_MULTI_OF)
            for cap_feat in all_cap_feats
        )

        if caption_valid_lens is not None:
            caption_valid_lens = caption_valid_lens.to(
                device=all_cap_feats[0].device, dtype=torch.long
            )

        for idx, cap_feat in enumerate(all_cap_feats):
            cap_ori_len = cap_feat.size(0)
            cap_attn_len = self._ceil_to_multiple(cap_ori_len, SEQ_MULTI_OF)
            cap_padding_len = cap_seq_len_target - cap_ori_len
            cap_padded_feat = torch.cat(
                [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)],
                dim=0,
            )
            all_cap_feats_out.append(cap_padded_feat)
            if caption_valid_mask is not None:
                mask_row = caption_valid_mask[idx].to(
                    device=cap_feat.device, dtype=torch.bool
                )
                if mask_row.dim() != 1:
                    mask_row = mask_row.reshape(-1)
                if mask_row.shape[0] > cap_seq_len_target:
                    mask_row = mask_row[:cap_seq_len_target]
                elif mask_row.shape[0] < cap_seq_len_target:
                    mask_row = torch.nn.functional.pad(
                        mask_row,
                        (0, cap_seq_len_target - mask_row.shape[0]),
                        value=0,
                    )
                all_cap_valid_masks.append(mask_row)
            if caption_valid_lens is None:
                all_cap_valid_lens.append(cap_ori_len)
            else:
                all_cap_valid_lens.append(caption_valid_lens[idx])
            all_cap_attn_lens.append(cap_attn_len)

        target_image_seq_len = image_seq_len_target or 0
        for image in all_image:
            # ------------ Process Image ------------
            C, F, H, W = image.size()
            image_size = (F, H, W)

            F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
            image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
            # "c f pf h ph w pw -> (f h w) (pf ph pw c)"
            image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(
                F_tokens * H_tokens * W_tokens, pF * pH * pW * C
            )
            image_ori_len = image.size(0)
            image_attn_len = max(
                image_seq_len_target or 0,
                self._ceil_to_multiple(image_ori_len, SEQ_MULTI_OF),
            )
            target_image_seq_len = max(
                target_image_seq_len,
                image_attn_len,
            )
            image_records.append((image, image_size, image_ori_len, image_attn_len))

        for image, image_size, image_ori_len, image_attn_len in image_records:
            image_padding_len = target_image_seq_len - image_ori_len
            image_padded_feat = torch.cat(
                [image, image[-1:].repeat(image_padding_len, 1)],
                dim=0,
            )
            all_image_out.append(image_padded_feat)
            all_image_size.append(image_size)
            all_image_valid_lens.append(image_ori_len)
            all_image_attn_lens.append(image_attn_len)

        cap_valid_lens_out = (
            caption_valid_lens if caption_valid_lens is not None else all_cap_valid_lens
        )
        return (
            torch.stack(all_image_out, dim=0),
            torch.stack(all_cap_feats_out, dim=0),
            all_image_size,
            all_image_valid_lens,
            cap_valid_lens_out,
            all_image_attn_lens,
            all_cap_attn_lens,
            (
                torch.stack(all_cap_valid_masks, dim=0)
                if caption_valid_mask is not None
                else None
            ),
        )

    def _build_single_sample_freqs_cis(
        self,
        image: torch.Tensor,
        cap_feat: torch.Tensor,
        patch_size: int,
        f_patch_size: int,
    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
        device = image.device
        cap_ori_len = int(cap_feat.size(0))
        cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
        cap_pos_ids = self.create_coordinate_grid(
            size=(cap_ori_len + cap_padding_len, 1, 1),
            start=(1, 0, 0),
            device=device,
        ).flatten(0, 2)

        _, F, H, W = image.size()
        pH = pW = patch_size
        pF = f_patch_size
        F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
        image_ori_len = F_tokens * H_tokens * W_tokens
        image_padding_len = (-image_ori_len) % SEQ_MULTI_OF
        image_ori_pos_ids = self.create_coordinate_grid(
            size=(F_tokens, H_tokens, W_tokens),
            start=(cap_ori_len + cap_padding_len + 1, 0, 0),
            device=device,
        ).flatten(0, 2)
        image_padding_pos_ids = (
            self.create_coordinate_grid(
                size=(1, 1, 1),
                start=(0, 0, 0),
                device=device,
            )
            .flatten(0, 2)
            .repeat(image_padding_len, 1)
        )
        image_pos_ids = torch.cat([image_ori_pos_ids, image_padding_pos_ids], dim=0)

        return self.rotary_emb(cap_pos_ids), self.rotary_emb(image_pos_ids)

    @staticmethod
    def _pad_freqs_cis_to_length(
        freqs_cis: Tuple[torch.Tensor, torch.Tensor], target_len: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        cos, sin = freqs_cis
        pad_len = target_len - cos.shape[0]
        if pad_len < 0:
            raise ValueError(
                f"Cannot pad RoPE freqs of length {cos.shape[0]} to shorter target {target_len}"
            )
        if pad_len == 0:
            return cos, sin
        return (
            torch.cat([cos, cos.new_zeros(pad_len, cos.shape[-1])], dim=0),
            torch.cat([sin, sin.new_zeros(pad_len, sin.shape[-1])], dim=0),
        )

    def _build_batched_freqs_cis(
        self,
        images: list[torch.Tensor],
        cap_feats: list[torch.Tensor],
        patch_size: int,
        f_patch_size: int,
        image_target_len: int,
        cap_target_len: int,
    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
        cap_cos, cap_sin, image_cos, image_sin = [], [], [], []
        for image, cap_feat in zip(images, cap_feats):
            sample_cap_freqs, sample_image_freqs = self._build_single_sample_freqs_cis(
                image,
                cap_feat,
                patch_size,
                f_patch_size,
            )
            sample_cap_freqs = self._pad_freqs_cis_to_length(
                sample_cap_freqs, cap_target_len
            )
            sample_image_freqs = self._pad_freqs_cis_to_length(
                sample_image_freqs, image_target_len
            )
            cap_cos.append(sample_cap_freqs[0])
            cap_sin.append(sample_cap_freqs[1])
            image_cos.append(sample_image_freqs[0])
            image_sin.append(sample_image_freqs[1])

        return (
            (torch.stack(cap_cos, dim=0), torch.stack(cap_sin, dim=0)),
            (torch.stack(image_cos, dim=0), torch.stack(image_sin, dim=0)),
        )

    @staticmethod
    def _device_cache_key(device: torch.device) -> tuple[str, int | None]:
        device = torch.device(device)
        return device.type, device.index

    def _get_cached_batched_freqs_cis(
        self,
        images: list[torch.Tensor],
        cap_feats: list[torch.Tensor],
        patch_size: int,
        f_patch_size: int,
        image_target_len: int,
        cap_target_len: int,
        device: torch.device,
    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
        cache_key = (
            tuple(tuple(image.shape) for image in images),
            tuple(tuple(cap_feat.shape) for cap_feat in cap_feats),
            int(patch_size),
            int(f_patch_size),
            int(image_target_len),
            int(cap_target_len),
            self._device_cache_key(device),
        )
        cached = getattr(self, "_cached_batched_freqs_cis", None)
        if cached is not None and cached[0] == cache_key:
            return cached[1]

        freqs_cis = self._build_batched_freqs_cis(
            images,
            cap_feats,
            patch_size,
            f_patch_size,
            image_target_len=image_target_len,
            cap_target_len=cap_target_len,
        )
        self._cached_batched_freqs_cis = (cache_key, freqs_cis)
        return freqs_cis

    def _get_rope_cache(
        self,
        cache_attr: str,
        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]],
    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
        if freqs_cis is None or not _is_cuda:
            return None, None

        cos, sin = freqs_cis
        if not (cos.is_cuda and sin.is_cuda):
            return None, None

        cache_key = (
            cos.data_ptr(),
            sin.data_ptr(),
            tuple(cos.shape),
            tuple(sin.shape),
            cos.dtype,
            sin.dtype,
            self._device_cache_key(cos.device),
        )
        cached = getattr(self, cache_attr, None)
        if cached is not None and cached[0] == cache_key:
            return cached[1]

        if cos.dim() == 3:
            batch_size, seq_len = cos.shape[:2]
            cos_sin_cache = torch.cat(
                [
                    cos.to(dtype=torch.float32).contiguous(),
                    sin.to(dtype=torch.float32).contiguous(),
                ],
                dim=-1,
            ).reshape(batch_size * seq_len, -1)
            positions = torch.arange(
                batch_size * seq_len, device=cos.device, dtype=torch.long
            )
        elif cos.dim() == 2:
            cos_sin_cache = torch.cat(
                [
                    cos.to(dtype=torch.float32).contiguous(),
                    sin.to(dtype=torch.float32).contiguous(),
                ],
                dim=-1,
            )
            positions = None
        else:
            return None, None

        rope_cache = (cos_sin_cache, positions)
        setattr(self, cache_attr, (cache_key, rope_cache))
        return rope_cache

    def _get_attn_mask_and_meta(
        self, cache_attr: str, lengths: list[int], target_len: int, device: torch.device
    ) -> Tuple[Optional[torch.Tensor], Optional[dict]]:
        length_key = tuple(int(length) for length in lengths)
        if all(length == target_len for length in length_key):
            return None, None

        cache_key = (
            length_key,
            int(target_len),
            self._device_cache_key(device),
        )
        cached = getattr(self, cache_attr, None)
        if cached is not None and cached[0] == cache_key:
            return cached[1]

        positions = torch.arange(target_len, device=device).unsqueeze(0)
        length_tensor = torch.as_tensor(
            length_key, dtype=torch.long, device=device
        ).unsqueeze(1)
        mask = positions < length_tensor
        meta = build_varlen_mask_meta_from_lengths(length_key, target_len, device)
        result = (mask, meta)
        setattr(self, cache_attr, (cache_key, result))
        return result

    def _get_joint_attn_mask_and_meta(
        self,
        image_lengths: list[int],
        image_target_len: int,
        cap_lengths: list[int],
        cap_target_len: int,
        device: torch.device,
    ) -> Tuple[Optional[torch.Tensor], Optional[dict]]:
        image_length_key = tuple(int(length) for length in image_lengths)
        cap_length_key = tuple(int(length) for length in cap_lengths)
        if all(length == image_target_len for length in image_length_key) and all(
            length == cap_target_len for length in cap_length_key
        ):
            return None, None

        cache_key = (
            image_length_key,
            int(image_target_len),
            cap_length_key,
            int(cap_target_len),
            self._device_cache_key(device),
        )
        cached = getattr(self, "_cached_joint_attn_mask_meta", None)
        if cached is not None and cached[0] == cache_key:
            return cached[1]

        image_pos = torch.arange(image_target_len, device=device).unsqueeze(0)
        cap_pos = torch.arange(cap_target_len, device=device).unsqueeze(0)
        image_len = torch.as_tensor(
            image_length_key, dtype=torch.long, device=device
        ).unsqueeze(1)
        cap_len = torch.as_tensor(
            cap_length_key, dtype=torch.long, device=device
        ).unsqueeze(1)
        mask = torch.cat([image_pos < image_len, cap_pos < cap_len], dim=1)
        valid_ranges = [
            [
                (0, image_length),
                (image_target_len, image_target_len + cap_length),
            ]
            for image_length, cap_length in zip(
                image_length_key, cap_length_key, strict=True
            )
        ]
        meta = build_varlen_mask_meta_from_ranges(
            valid_ranges,
            image_target_len + cap_target_len,
            device,
        )
        result = (mask, meta)
        self._cached_joint_attn_mask_meta = (cache_key, result)
        return result

    @staticmethod
    def _has_padding(valid_lens: list[int], target_len: int) -> bool:
        return any(int(length) < target_len for length in valid_lens)

    @staticmethod
    def _as_image_list(hidden_states) -> list[torch.Tensor]:
        """Normalize 4D/5D image latents into per-sample tensors."""
        if torch.is_tensor(hidden_states):
            if hidden_states.dim() == 5:
                return list(hidden_states.unbind(dim=0))
            if hidden_states.dim() == 4:
                return [hidden_states]
        return list(hidden_states)

    @staticmethod
    def _as_caption_list(encoder_hidden_states) -> list[torch.Tensor]:
        """Normalize caption tensors into per-sample tensors."""
        if torch.is_tensor(encoder_hidden_states):
            if encoder_hidden_states.dim() == 3:
                return list(encoder_hidden_states.unbind(dim=0))
            if encoder_hidden_states.dim() == 2:
                return [encoder_hidden_states]

        cap_feats = list(encoder_hidden_states)
        if len(cap_feats) == 1 and torch.is_tensor(cap_feats[0]):
            if cap_feats[0].dim() == 3:
                return list(cap_feats[0].unbind(dim=0))
            if cap_feats[0].dim() == 2:
                return cap_feats
        return cap_feats

    @staticmethod
    def _caption_valid_mask_from_mask(
        mask, *, batch_size: int, max_seq_len: int
    ) -> torch.Tensor | None:
        if mask is None:
            return None
        if isinstance(mask, (list, tuple)):
            if not mask:
                return None
            if len(mask) == 1:
                return ZImageTransformer2DModel._caption_valid_mask_from_mask(
                    mask[0], batch_size=batch_size, max_seq_len=max_seq_len
                )
            rows = []
            for item in mask:
                item_mask = ZImageTransformer2DModel._caption_valid_mask_from_mask(
                    item, batch_size=1, max_seq_len=max_seq_len
                )
                if item_mask is None:
                    return None
                rows.append(item_mask[0])
            return torch.stack(rows, dim=0) if len(rows) == batch_size else None
        if not torch.is_tensor(mask):
            return None

        mask = mask.to(dtype=torch.bool)
        if mask.ndim == 1:
            if batch_size != 1:
                return None
            mask = mask[:max_seq_len].unsqueeze(0)
        elif mask.ndim == 2 and mask.shape[0] == batch_size:
            mask = mask[:, :max_seq_len]
        elif mask.ndim == 2 and batch_size == 1 and mask.shape[0] == 1:
            mask = mask[:, :max_seq_len]
        else:
            return None

        return mask

    @staticmethod
    def _replace_padding_with_token(
        tensor: torch.Tensor,
        valid_lens: list[int] | torch.Tensor,
        pad_token: torch.Tensor,
    ) -> torch.Tensor:
        """Replace padded token rows after each valid sequence length."""
        if not torch.is_tensor(valid_lens) and all(
            int(length) >= tensor.shape[1] for length in valid_lens
        ):
            return tensor
        positions = torch.arange(tensor.shape[1], device=tensor.device).unsqueeze(0)
        if torch.is_tensor(valid_lens):
            lengths = valid_lens.to(device=tensor.device, dtype=torch.long)
        else:
            lengths = torch.tensor(valid_lens, device=tensor.device)
        if lengths.ndim == 0:
            lengths = lengths.reshape(1)
        lengths = lengths.unsqueeze(1)
        pad_mask = positions >= lengths
        tensor = tensor.clone()
        tensor[pad_mask] = pad_token.to(device=tensor.device, dtype=tensor.dtype)
        return tensor

    @staticmethod
    def _replace_padding_with_token_mask(
        tensor: torch.Tensor,
        valid_mask: torch.Tensor,
        pad_token: torch.Tensor,
    ) -> torch.Tensor:
        """Replace padded token rows using a fixed-shape tensor mask."""
        seq_len = tensor.shape[1]
        valid_mask = valid_mask.to(device=tensor.device, dtype=torch.bool)
        if valid_mask.shape[1] > seq_len:
            valid_mask = valid_mask[:, :seq_len]
        elif valid_mask.shape[1] < seq_len:
            valid_mask = torch.nn.functional.pad(
                valid_mask,
                (0, seq_len - valid_mask.shape[1]),
                value=0,
            )
        pad_value = pad_token.to(device=tensor.device, dtype=tensor.dtype)
        return torch.where(valid_mask.unsqueeze(-1), tensor, pad_value.view(1, 1, -1))

    def forward(
        self,
        hidden_states: List[torch.Tensor],
        encoder_hidden_states: List[torch.Tensor],
        timestep,
        guidance=0,
        patch_size=2,
        f_patch_size=1,
        freqs_cis=None,
        image_seq_len_target: int | None = None,
        encoder_hidden_states_mask=None,
        caption_valid_lens: torch.Tensor | None = None,
        **kwargs,
    ):
        assert patch_size in self.all_patch_size
        assert f_patch_size in self.all_f_patch_size

        x = self._as_image_list(hidden_states)
        cap_feats = self._as_caption_list(encoder_hidden_states)
        input_images = x
        input_cap_feats = cap_feats
        caption_valid_mask = None
        if kwargs.pop("_use_caption_valid_mask", False):
            caption_valid_mask = self._caption_valid_mask_from_mask(
                encoder_hidden_states_mask,
                batch_size=len(cap_feats),
                max_seq_len=max(cap_feat.shape[0] for cap_feat in cap_feats),
            )
        timestep = 1000.0 - timestep
        t = timestep
        t = self.t_embedder(t)
        adaln_input = t.to(dtype=x[0].dtype)
        (
            x,
            cap_feats,
            x_size,
            x_valid_lens,
            cap_valid_lens,
            x_attn_lens,
            cap_attn_lens,
            cap_valid_mask,
        ) = self.patchify_and_embed(
            x,
            cap_feats,
            patch_size,
            f_patch_size,
            image_seq_len_target=image_seq_len_target,
            caption_valid_lens=caption_valid_lens,
            caption_valid_mask=caption_valid_mask,
        )

        x, _ = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x)
        device = x.device
        x = self._replace_padding_with_token(x, x_valid_lens, self.x_pad_token)
        if len(input_images) > 1 and get_sp_world_size() == 1:
            freqs_cis = self._get_cached_batched_freqs_cis(
                input_images,
                input_cap_feats,
                patch_size,
                f_patch_size,
                image_target_len=x.shape[1],
                cap_target_len=cap_feats.shape[1],
                device=device,
            )
        x_freqs_cis = freqs_cis[1]
        x_rope_cos_sin_cache, x_rope_positions = self._get_rope_cache(
            "_cached_x_rope_cache", x_freqs_cis
        )
        x_attn_mask, x_attn_mask_meta = self._get_attn_mask_and_meta(
            "_cached_x_attn_mask_meta", x_attn_lens, x.shape[1], device
        )

        for layer_id, layer in enumerate(self.noise_refiner):
            x = layer(
                x,
                x_freqs_cis,
                adaln_input,
                rope_cos_sin_cache=x_rope_cos_sin_cache,
                rope_positions=x_rope_positions,
                attn_mask=x_attn_mask,
                attn_mask_meta=x_attn_mask_meta,
            )

        cap_feats, _ = self.cap_embedder(cap_feats)
        if cap_valid_mask is not None:
            cap_feats = self._replace_padding_with_token_mask(
                cap_feats, cap_valid_mask, self.cap_pad_token
            )
        else:
            cap_feats = self._replace_padding_with_token(
                cap_feats, cap_valid_lens, self.cap_pad_token
            )

        cap_freqs_cis = freqs_cis[0]
        cap_rope_cos_sin_cache, cap_rope_positions = self._get_rope_cache(
            "_cached_cap_rope_cache", cap_freqs_cis
        )
        cap_attn_mask, cap_attn_mask_meta = self._get_attn_mask_and_meta(
            "_cached_cap_attn_mask_meta", cap_attn_lens, cap_feats.shape[1], device
        )

        for layer_id, layer in enumerate(self.context_refiner):
            cap_feats = layer(
                cap_feats,
                cap_freqs_cis,
                rope_cos_sin_cache=cap_rope_cos_sin_cache,
                rope_positions=cap_rope_positions,
                attn_mask=cap_attn_mask,
                attn_mask_meta=cap_attn_mask_meta,
            )

        cap_seq_len = cap_feats.shape[1]
        use_full_unified_sequence = (
            get_sp_world_size() > 1 and get_ring_parallel_world_size() > 1
        )
        x_local_seq_len = x.shape[1]
        if use_full_unified_sequence:
            x = sequence_model_parallel_all_gather(x.contiguous(), dim=1)
            x_freqs_cis = (
                sequence_model_parallel_all_gather(x_freqs_cis[0].contiguous(), dim=0),
                sequence_model_parallel_all_gather(x_freqs_cis[1].contiguous(), dim=0),
            )
        unified = torch.cat([x, cap_feats], dim=1)
        unified_freqs_cis = (
            torch.cat([x_freqs_cis[0], cap_freqs_cis[0]], dim=-2),
            torch.cat([x_freqs_cis[1], cap_freqs_cis[1]], dim=-2),
        )
        unified_attn_mask, unified_attn_mask_meta = self._get_joint_attn_mask_and_meta(
            x_attn_lens,
            x.shape[1],
            cap_attn_lens,
            cap_seq_len,
            device,
        )
        unified_rope_cos_sin_cache, unified_rope_positions = self._get_rope_cache(
            "_cached_unified_rope_cache", unified_freqs_cis
        )
        num_replicated_suffix = cap_seq_len if not use_full_unified_sequence else 0

        for layer in self.layers:
            unified = layer(
                unified,
                unified_freqs_cis,
                adaln_input,
                rope_cos_sin_cache=unified_rope_cos_sin_cache,
                rope_positions=unified_rope_positions,
                attn_mask=unified_attn_mask,
                attn_mask_meta=unified_attn_mask_meta,
                num_replicated_suffix=num_replicated_suffix,
                skip_sequence_parallel_override=use_full_unified_sequence,
            )

        unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](
            unified, adaln_input
        )
        if use_full_unified_sequence:
            sp_rank = get_sp_parallel_rank()
            start = sp_rank * x_local_seq_len
            end = start + x_local_seq_len
            unified = unified[:, start:end]
        x = list(unified.unbind(dim=0))
        x = self.unpatchify(x, x_size, patch_size, f_patch_size)

        # Keep batch dim so output shape matches input (e.g. rollout/scheduler expect same ndim).
        return -torch.stack(x)


EntryClass = ZImageTransformer2DModel