[diffusion] fix: fix z-Image accuracy (#29742)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/jit_kernel/diffusion/triton/zimage_native_norm.py

@@ -0,0 +1,182 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+
+@triton.jit
+def _tanh(x):
+    return 2.0 / (1.0 + tl.exp(-2.0 * x)) - 1.0
+
+
+@triton.jit
+def _rmsnorm_scale_kernel(
+    y_ptr,
+    x_ptr,
+    weight_ptr,
+    scale_ptr,
+    x_row_stride,
+    scale_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    scale = tl.load(
+        scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0
+    )
+    y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+@triton.jit
+def _rmsnorm_tanh_residual_kernel(
+    y_ptr,
+    x_ptr,
+    gate_ptr,
+    residual_ptr,
+    weight_ptr,
+    x_row_stride,
+    gate_row_stride,
+    residual_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    gate = tl.load(gate_ptr + batch * gate_row_stride + offsets, mask=mask, other=0.0)
+    residual = tl.load(
+        residual_ptr + row * residual_row_stride + offsets, mask=mask, other=0.0
+    )
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    norm = ((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16)
+    gated = (_tanh(gate.to(tl.float32)).to(tl.bfloat16) * norm).to(tl.bfloat16)
+    y = (residual + gated).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+def _flat_row_stride(x: torch.Tensor) -> int | None:
+    if x.dim() < 2 or x.stride(-1) != 1:
+        return None
+    row_stride = x.stride(-2)
+    expected_stride = row_stride * x.shape[-2]
+    for dim in range(x.dim() - 3, -1, -1):
+        if x.stride(dim) != expected_stride:
+            return None
+        expected_stride *= x.shape[dim]
+    return row_stride
+
+
+def _can_use(x: torch.Tensor, weight: torch.Tensor, other: torch.Tensor) -> bool:
+    return (
+        x.is_cuda
+        and weight.is_cuda
+        and other.is_cuda
+        and x.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and other.dtype == torch.bfloat16
+        and weight.is_contiguous()
+        and x.shape[-1] <= 8192
+        and _flat_row_stride(x) is not None
+        and _flat_row_stride(other) is not None
+    )
+
+
+def zimage_rmsnorm_scale(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    if not _can_use(x, weight, scale):
+        return None
+    shape = x.shape
+    dim = shape[-1]
+    x_rows = x.numel() // dim
+    scale_rows = scale.numel() // dim
+    if x_rows % scale_rows != 0:
+        return None
+    seq_len = x_rows // scale_rows
+    x_row_stride = _flat_row_stride(x)
+    scale_row_stride = _flat_row_stride(scale)
+    if x_row_stride is None or scale_row_stride is None:
+        return None
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_scale_kernel[(x_rows,)](
+            y.reshape(-1, dim),
+            x,
+            weight,
+            scale,
+            x_row_stride,
+            scale_row_stride,
+            seq_len,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return y
+
+
+def zimage_rmsnorm_tanh_residual(
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    if not (_can_use(x, weight, gate) and residual.is_cuda):
+        return None
+    if residual.dtype != x.dtype or _flat_row_stride(residual) is None:
+        return None
+    shape = x.shape
+    dim = shape[-1]
+    x_rows = x.numel() // dim
+    gate_rows = gate.numel() // dim
+    if x_rows % gate_rows != 0:
+        return None
+    seq_len = x_rows // gate_rows
+    x_row_stride = _flat_row_stride(x)
+    gate_row_stride = _flat_row_stride(gate)
+    residual_row_stride = _flat_row_stride(residual)
+    if x_row_stride is None or gate_row_stride is None or residual_row_stride is None:
+        return None
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_tanh_residual_kernel[(x_rows,)](
+            y.reshape(-1, dim),
+            x,
+            gate,
+            residual,
+            weight,
+            x_row_stride,
+            gate_row_stride,
+            residual_row_stride,
+            seq_len,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return y

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

@@ -67,6 +67,7 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
     task_type: ModelTaskType = ModelTaskType.T2I
     dit_config: DiTConfig = field(default_factory=ZImageDitConfig)
     vae_config: VAEConfig = field(default_factory=FluxVAEConfig)
+    enable_autocast: bool = False
     vae_precision: str = "bf16"
     text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
     text_encoder_configs: tuple[EncoderConfig, ...] = field(
@@ -87,6 +88,9 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(fsdp_auto_min_available_memory_gb=40)
 
+    def prepare_sigmas(self, sigmas, num_inference_steps):
+        return self._prepare_sigmas(sigmas, num_inference_steps)
+
     def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
         rendered_prompts = [
             tokenizer.apply_chat_template(

python/sglang/multimodal_gen/runtime/layers/attention/__init__.py

@@ -13,6 +13,8 @@ from sglang.multimodal_gen.runtime.layers.attention.layer import (
     UlyssesAttention_VSA,
     USPAttention,
     build_varlen_mask_meta,
+    build_varlen_mask_meta_from_lengths,
+    build_varlen_mask_meta_from_ranges,
 )
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
 from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import MinimalA2AAttnOp
@@ -29,4 +31,6 @@ __all__ = [
     # "AttentionState",
     "get_attn_backend",
     "build_varlen_mask_meta",
+    "build_varlen_mask_meta_from_lengths",
+    "build_varlen_mask_meta_from_ranges",
 ]

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -2,6 +2,7 @@
 
 # SPDX-License-Identifier: Apache-2.0
 import os
+from collections.abc import Sequence
 from contextlib import nullcontext
 from typing import Type
 
@@ -92,6 +93,84 @@ def build_varlen_mask_meta(
     }
 
 
+def build_varlen_mask_meta_from_lengths(
+    lengths: Sequence[int],
+    max_seqlen: int,
+    device: torch.device,
+) -> dict:
+    """Build varlen FA metadata for prefix-valid masks without CUDA nonzero.
+
+    This is equivalent to ``build_varlen_mask_meta`` for masks where row ``i`` is
+    true on ``[:lengths[i]]`` and false afterwards. Keeping the lengths on the
+    host lets callers avoid a GPU ``nonzero``/dynamic-shape path while still
+    producing the same packed indices.
+    """
+
+    return build_varlen_mask_meta_from_ranges(
+        [[(0, int(length))] for length in lengths],
+        max_seqlen=max_seqlen,
+        device=device,
+    )
+
+
+def build_varlen_mask_meta_from_ranges(
+    valid_ranges: Sequence[Sequence[tuple[int, int]]],
+    max_seqlen: int,
+    device: torch.device,
+) -> dict:
+    """Build varlen FA metadata from host-side valid token ranges.
+
+    ``valid_ranges[i]`` contains half-open intervals in row-local coordinates.
+    The intervals are packed in the provided order, matching the flattened
+    ``nonzero`` order for ordinary left-to-right masks.
+    """
+
+    range_values = [
+        [(int(start), int(end)) for start, end in row_ranges]
+        for row_ranges in valid_ranges
+    ]
+    if any(
+        start < 0 or end < start or end > max_seqlen
+        for row_ranges in range_values
+        for start, end in row_ranges
+    ):
+        raise ValueError(
+            f"All ranges must be within [0, {max_seqlen}], got {range_values}"
+        )
+
+    bs = len(range_values)
+    length_values = [
+        sum(end - start for start, end in row_ranges) for row_ranges in range_values
+    ]
+    valid_lens = torch.as_tensor(length_values, dtype=torch.int32, device=device)
+    cu_seqlens = torch.zeros(bs + 1, dtype=torch.int32, device=device)
+    cu_seqlens[1:] = torch.cumsum(valid_lens, dim=0)
+
+    index_parts = [
+        torch.arange(
+            row * max_seqlen + start,
+            row * max_seqlen + end,
+            dtype=torch.long,
+            device=device,
+        )
+        for row, row_ranges in enumerate(range_values)
+        for start, end in row_ranges
+        if end > start
+    ]
+    if index_parts:
+        indices = torch.cat(index_parts, dim=0)
+    else:
+        indices = torch.empty((0,), dtype=torch.long, device=device)
+    inv_indices = build_inv_indices(indices, bs * max_seqlen)
+
+    return {
+        "cu_seqlens": cu_seqlens,
+        "indices": indices,
+        "inv_indices": inv_indices,
+        "max_seqlen": max_seqlen,
+    }
+
+
 class UlyssesAttention(nn.Module):
     """Ulysses-style SequenceParallelism attention layer."""
 

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

@@ -18,11 +18,11 @@ 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 (
-    RMSNorm,
     apply_qk_norm_with_optional_rope,
-    apply_rmsnorm_tanh_mul_add,
 )
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
@@ -60,6 +60,73 @@ 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__()
@@ -167,6 +234,7 @@ class ZImageAttention(nn.Module):
         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 (
@@ -217,8 +285,8 @@ class ZImageAttention(nn.Module):
             )
 
         if self.qk_norm:
-            self.norm_q = RMSNorm(self.head_dim, eps=eps)
+            self.norm_q = ZImageRMSNorm(self.head_dim, eps=eps)
-            self.norm_k = RMSNorm(self.head_dim, eps=eps)
+            self.norm_k = ZImageRMSNorm(self.head_dim, eps=eps)
         else:
             self.norm_q = None
             self.norm_k = None
@@ -256,6 +324,10 @@ class ZImageAttention(nn.Module):
         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,
@@ -281,9 +353,58 @@ class ZImageAttention(nn.Module):
         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 freqs_cis is not None:
+        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 _is_cuda and q.shape == k.shape:
+            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(),
@@ -300,7 +421,7 @@ class ZImageAttention(nn.Module):
                         head_dim=self.head_dim,
                         cos_sin_cache=cos_sin_cache,
                         is_neox=False,
-                        allow_inplace=True,
+                        allow_inplace=self.enable_zimage_qk_fusion,
                     )
                 else:
                     q, k = apply_flashinfer_rope_qk_inplace(
@@ -314,7 +435,7 @@ class ZImageAttention(nn.Module):
                         q_norm=self.norm_q,
                         k_norm=self.norm_k,
                         head_dim=self.head_dim,
-                        allow_inplace=True,
+                        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)
@@ -325,7 +446,7 @@ class ZImageAttention(nn.Module):
                 q_norm=self.norm_q,
                 k_norm=self.norm_k,
                 head_dim=self.head_dim,
-                allow_inplace=True,
+                allow_inplace=self.enable_zimage_qk_fusion,
             )
 
         if (
@@ -361,6 +482,8 @@ class ZImageAttention(nn.Module):
                 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,
@@ -390,6 +513,7 @@ class ZImageTransformerBlock(nn.Module):
         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,
@@ -438,11 +562,11 @@ class ZImageTransformerBlock(nn.Module):
                 prefix=f"{prefix}.feed_forward",
             )
 
-        self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.attention_norm1 = ZImageRMSNorm(dim, eps=norm_eps)
-        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm1 = ZImageRMSNorm(dim, eps=norm_eps)
 
-        self.attention_norm2 = RMSNorm(dim, eps=norm_eps)
+        self.attention_norm2 = ZImageRMSNorm(dim, eps=norm_eps)
-        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = ZImageRMSNorm(dim, eps=norm_eps)
 
         if modulation:
             self.adaLN_modulation = nn.Sequential(
@@ -454,6 +578,10 @@ class ZImageTransformerBlock(nn.Module):
         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,
@@ -468,51 +596,54 @@ class ZImageTransformerBlock(nn.Module):
 
             # Attention block
             attn_out = self.attention(
-                self.attention_norm1(x) * scale_msa,
+                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,
             )
-            if (
+            x = zimage_rmsnorm_tanh_mul_add(
-                _is_cuda
+                attn_out,
-                and attn_out.is_cuda
+                gate_msa,
-                and attn_out.shape[-1] % 256 == 0
+                x,
-                and attn_out.shape[-1] <= 8192
+                self.attention_norm2,
-                and self.attention_norm2.variance_epsilon
+                self.enable_zimage_native_norm_fusion,
-                == self.ffn_norm1.variance_epsilon
+            )
-            ):
+            ffn_in = zimage_rmsnorm_scale(
-                from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
+                x,
-                    fused_norm_tanh_mul_add_norm_scale,
+                1.0 + scale_mlp,
-                )
+                self.ffn_norm1,
-
+                self.enable_zimage_native_norm_fusion,
-                x, ffn_in = fused_norm_tanh_mul_add_norm_scale(
+            )
-                    attn_out.contiguous(),
-                    self.attention_norm2.weight.data.contiguous(),
-                    None,
-                    gate_msa.contiguous(),
-                    x.contiguous(),
-                    self.ffn_norm1.weight.data.contiguous(),
-                    None,
-                    scale_mlp.contiguous(),
-                    "rms",
-                    self.attention_norm2.variance_epsilon,
-                )
-            else:
-                x = apply_rmsnorm_tanh_mul_add(
-                    attn_out, gate_msa, x, self.attention_norm2
-                )
-                ffn_in = self.ffn_norm1(x) * (1.0 + scale_mlp)
 
             # FFN block
             ffn_out = self.feed_forward(ffn_in)
-            x = apply_rmsnorm_tanh_mul_add(ffn_out, gate_mlp, x, self.ffn_norm2)
+            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,
@@ -741,7 +872,7 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         )
 
         self.cap_embedder = nn.Sequential(
-            RMSNorm(arch_config.cap_feat_dim, eps=arch_config.norm_eps),
+            ZImageRMSNorm(arch_config.cap_feat_dim, eps=arch_config.norm_eps),
             ReplicatedLinear(arch_config.cap_feat_dim, self.dim, bias=True),
         )
 
@@ -840,6 +971,8 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         all_cap_feats_out = []
         all_image_valid_lens = []
         all_cap_valid_lens = []
+        all_image_attn_lens = []
+        all_cap_attn_lens = []
         image_records = []
 
         cap_seq_len_target = max(
@@ -854,6 +987,7 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
 
         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)],
@@ -864,6 +998,7 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
                 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:
@@ -878,13 +1013,17 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
                 F_tokens * H_tokens * W_tokens, pF * pH * pW * C
             )
             image_ori_len = image.size(0)
-            target_image_seq_len = max(
+            image_attn_len = max(
-                target_image_seq_len,
+                image_seq_len_target or 0,
                 self._ceil_to_multiple(image_ori_len, SEQ_MULTI_OF),
             )
-            image_records.append((image, image_size, image_ori_len))
+            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 in image_records:
+        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)],
@@ -893,6 +1032,7 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
             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
@@ -903,8 +1043,275 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
             all_image_size,
             all_image_valid_lens,
             cap_valid_lens_out,
+            all_image_attn_lens,
+            all_cap_attn_lens,
         )
 
+    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."""
@@ -939,6 +1346,9 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         pad_token: torch.Tensor,
     ) -> torch.Tensor:
         """Replace padded token rows after each valid sequence length."""
+        if not ZImageTransformer2DModel._has_padding(valid_lens, tensor.shape[1]):
+            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)
@@ -946,9 +1356,8 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
             lengths = torch.tensor(valid_lens, device=tensor.device)
         lengths = lengths.unsqueeze(1)
         pad_mask = positions >= lengths
-        if pad_mask.any():
+        tensor = tensor.clone()
-            tensor = tensor.clone()
+        tensor[pad_mask] = pad_token.to(device=tensor.device, dtype=tensor.dtype)
-            tensor[pad_mask] = pad_token.to(device=tensor.device, dtype=tensor.dtype)
         return tensor
 
     def forward(
@@ -969,6 +1378,9 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
 
         x = self._as_image_list(hidden_states)
         cap_feats = self._as_caption_list(encoder_hidden_states)
+        input_images = x
+        input_cap_feats = cap_feats
+
         timestep = 1000.0 - timestep
         t = timestep
         t = self.t_embedder(t)
@@ -979,6 +1391,8 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
             x_size,
             x_valid_lens,
             cap_valid_lens,
+            x_attn_lens,
+            cap_attn_lens,
         ) = self.patchify_and_embed(
             x,
             cap_feats,
@@ -989,11 +1403,36 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         )
 
         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)
+            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)
         cap_feats = self._replace_padding_with_token(
@@ -1001,11 +1440,21 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         )
 
         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]
@@ -1021,8 +1470,18 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
             )
         unified = torch.cat([x, cap_feats], dim=1)
         unified_freqs_cis = (
-            torch.cat([x_freqs_cis[0], cap_freqs_cis[0]], dim=0),
+            torch.cat([x_freqs_cis[0], cap_freqs_cis[0]], dim=-2),
-            torch.cat([x_freqs_cis[1], cap_freqs_cis[1]], dim=0),
+            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
 
@@ -1031,6 +1490,10 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
                 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,
             )

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

@@ -385,9 +385,11 @@ class Qwen3ForCausalLM(TextEncoder):
         residual = None
 
         if position_ids is None:
-            position_ids = torch.arange(
+            position_ids = (
-                0, hidden_states.shape[1], device=hidden_states.device
+                torch.arange(0, hidden_states.shape[1], device=hidden_states.device)
-            ).unsqueeze(0)
+                .unsqueeze(0)
+                .expand(hidden_states.shape[0], -1)
+            )
 
         attention_lengths = None
         if attention_mask is not None:

python/sglang/multimodal_gen/test/server/perf_baselines/h100.json

@@ -653,78 +653,78 @@
         },
         "zimage_image_t2i": {
             "stages_ms": {
-                "DecodingStage": 8.86,
+                "DecodingStage": 134.1,
                 "InputValidationStage": 0.04,
-                "DenoisingStage": 675.8,
+                "DenoisingStage": 771.7,
-                "TextEncodingStage": 130.38,
+                "TextEncodingStage": 130.68,
-                "LatentPreparationStage": 0.14,
+                "LatentPreparationStage": 0.16,
-                "TimestepPreparationStage": 31.38
+                "TimestepPreparationStage": 27.67
             },
             "denoise_step_ms": {
-                "0": 19.93,
+                "0": 95.97,
-                "1": 28.65,
+                "1": 29.8,
-                "2": 83.91,
+                "2": 91.06,
-                "3": 83.48,
+                "3": 91.41,
-                "4": 83.51,
+                "4": 93.26,
-                "5": 83.69,
+                "5": 90.11,
-                "6": 84.08,
+                "6": 90.88,
-                "7": 84.04,
+                "7": 91.93,
-                "8": 84.35
+                "8": 93.16
             },
-            "expected_e2e_ms": 1027.94,
+            "expected_e2e_ms": 1068.66,
-            "expected_avg_denoise_ms": 74.6,
+            "expected_avg_denoise_ms": 85.29,
-            "expected_median_denoise_ms": 86.63,
+            "expected_median_denoise_ms": 91.41,
-            "estimated_full_test_time_s": 121.1
+            "estimated_full_test_time_s": 116.3
         },
         "zimage_image_t2i_fp8": {
             "stages_ms": {
-                "TextEncodingStage": 129.84,
+                "InputValidationStage": 0.04,
-                "DenoisingStage": 634.42,
+                "TextEncodingStage": 131.15,
-                "InputValidationStage": 0.03,
+                "LatentPreparationStage": 0.18,
-                "LatentPreparationStage": 0.11,
+                "TimestepPreparationStage": 26.29,
-                "TimestepPreparationStage": 17.42,
+                "DenoisingStage": 959.66,
-                "DecodingStage": 9.46
+                "DecodingStage": 125.62
             },
             "denoise_step_ms": {
-                "0": 33.35,
+                "0": 115.4,
-                "1": 36.43,
+                "1": 56.32,
-                "2": 66.88,
+                "2": 111.77,
-                "3": 78.15,
+                "3": 110.02,
-                "4": 78.0,
+                "4": 111.43,
-                "5": 78.36,
+                "5": 111.32,
-                "6": 78.48,
+                "6": 110.77,
-                "7": 78.32,
+                "7": 110.53,
-                "8": 73.36
+                "8": 118.39
             },
-            "expected_e2e_ms": 958.32,
+            "expected_e2e_ms": 1247.73,
-            "expected_avg_denoise_ms": 70.04,
+            "expected_avg_denoise_ms": 106.22,
-            "expected_median_denoise_ms": 81.35,
+            "expected_median_denoise_ms": 111.32,
-            "estimated_full_test_time_s": 121.0
+            "estimated_full_test_time_s": 123.7
         },
         "zimage_image_t2i_multi_lora": {
             "stages_ms": {
-                "TimestepPreparationStage": 24.93,
+                "InputValidationStage": 0.04,
-                "DenoisingStage": 673.95,
+                "TextEncodingStage": 130.25,
-                "DecodingStage": 8.43,
+                "LatentPreparationStage": 0.12,
-                "LatentPreparationStage": 0.1,
+                "TimestepPreparationStage": 26.5,
-                "TextEncodingStage": 129.4,
+                "DenoisingStage": 841.93,
-                "InputValidationStage": 0.03
+                "DecodingStage": 142.78
             },
             "denoise_step_ms": {
-                "0": 25.08,
+                "0": 102.18,
-                "1": 35.42,
+                "1": 40.06,
-                "2": 77.56,
+                "2": 99.39,
-                "3": 83.39,
+                "3": 100.33,
-                "4": 82.18,
+                "4": 103.95,
-                "5": 83.1,
+                "5": 93.48,
-                "6": 82.39,
+                "6": 99.6,
-                "7": 83.34,
+                "7": 99.38,
-                "8": 85.82
+                "8": 99.77
             },
-            "expected_e2e_ms": 1047.82,
+            "expected_e2e_ms": 1145.53,
-            "expected_avg_denoise_ms": 74.33,
+            "expected_avg_denoise_ms": 93.12,
-            "expected_median_denoise_ms": 86.33,
+            "expected_median_denoise_ms": 99.6,
-            "estimated_full_test_time_s": 121.1
+            "estimated_full_test_time_s": 162.1
         },
         "zimage_image_t2i_2_gpus": {
             "stages_ms": {

python/sglang/multimodal_gen/test/test_utils.py

@@ -34,7 +34,7 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-SGL_TEST_FILES_CI_DATA_REVISION = "916cbff23aa4e89f78128397ede7ce29a73d6d8c"
+SGL_TEST_FILES_CI_DATA_REVISION = "77bd016251220fee8917a30ec92e89da03794a8a"
 
 if current_platform.is_npu():
     SGL_TEST_FILES_CI_DATA_REVISION = "6b62f4b6825c76a25fd2ba28248df68f2b400e65"
@@ -96,6 +96,9 @@ DEFAULT_MEAN_ABS_DIFF_THRESHOLD_VIDEO = 10.0
 _clip_model_cache: dict[str, Any] = {}
 _consistency_gt_cache: dict[str, Any] = {}
 _official_consistency_gt_outputs_cache: dict[str, frozenset[str]] | None = None
+CONSISTENCY_GT_CASE_ALIASES = {
+    "fsdp-inference": "zimage_image_t2i_2_gpus",
+}
 OFFICIAL_CONSISTENCY_GT_SKIP_CASES = frozenset(
     {
         # Official references for these cases need regeneration or parity triage.
@@ -1016,10 +1019,15 @@ def output_format_to_ext(output_format: str | None) -> str:
     return "png"
 
 
+def get_consistency_gt_case_id(case_id: str) -> str:
+    return CONSISTENCY_GT_CASE_ALIASES.get(case_id, case_id)
+
+
 def _consistency_gt_filenames(
     case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None
 ) -> list[str]:
     """Return the list of GT image filenames for a case. Reused by GT generation and consistency check."""
+    case_id = get_consistency_gt_case_id(case_id)
     n = num_gpus
     if is_video:
         return [
@@ -1034,6 +1042,7 @@ def _consistency_gt_filenames(
 def _base_consistency_gt_candidates(
     case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None
 ) -> list[str]:
+    case_id = get_consistency_gt_case_id(case_id)
     n = num_gpus
     if is_video:
         return [
@@ -1053,9 +1062,9 @@ def get_consistency_gt_candidate_sets(
     candidates = _base_consistency_gt_candidates(
         case_id, num_gpus, is_video, output_format
     )
-    platform = get_consistency_platform()
+    if _is_ascend_consistency_case(case_id) or current_platform.is_npu():
-    if platform == "h100":
         return [candidates]
+    platform = get_consistency_platform()
     return [[f"{platform}/{candidate}" for candidate in candidates], candidates]
 
 
@@ -1157,7 +1166,8 @@ def _is_official_consistency_gt_base_url(base_url: str) -> bool:
 
 
 def _official_consistency_gt_candidate_is_declared(case_id: str, filename: str) -> bool:
-    return filename in _official_consistency_gt_outputs_for_case(case_id)
+    outputs = _official_consistency_gt_outputs_for_case(case_id)
+    return filename in outputs or filename.rsplit("/", 1)[-1] in outputs
 
 
 def _remote_consistency_gt_base_urls(case_id: str) -> tuple[str, ...]:

python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py

@@ -62,9 +62,9 @@ def test_remote_video_gt_candidates_survive_inconclusive_probe(monkeypatch):
     )
 
     assert [filename for filename, _ in files] == [
-        "unit_video_1gpu_frame_0.png",
+        "h100/unit_video_1gpu_frame_0.png",
-        "unit_video_1gpu_frame_mid.png",
+        "h100/unit_video_1gpu_frame_mid.png",
-        "unit_video_1gpu_frame_last.png",
+        "h100/unit_video_1gpu_frame_last.png",
     ]
 
 
@@ -87,10 +87,10 @@ def test_remote_image_gt_prefers_official_when_present(monkeypatch):
 
     assert files == [
         (
-            expected_filename,
+            f"h100/{expected_filename}",
             (
                 f"{test_utils.SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_BASE}"
-                f"/{expected_filename}"
+                f"/h100/{expected_filename}"
             ),
         )
     ]
@@ -116,10 +116,10 @@ def test_remote_image_gt_ignores_unmapped_official_file(monkeypatch):
 
     assert files == [
         (
-            expected_filename,
+            f"h100/{expected_filename}",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{expected_filename}"
+                f"/h100/{expected_filename}"
             ),
         )
     ]
@@ -141,24 +141,24 @@ def test_remote_video_gt_ignores_unmapped_official_files(monkeypatch):
 
     assert files == [
         (
-            f"{case_id}_2gpu_frame_0.png",
+            f"h100/{case_id}_2gpu_frame_0.png",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{case_id}_2gpu_frame_0.png"
+                f"/h100/{case_id}_2gpu_frame_0.png"
             ),
         ),
         (
-            f"{case_id}_2gpu_frame_mid.png",
+            f"h100/{case_id}_2gpu_frame_mid.png",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{case_id}_2gpu_frame_mid.png"
+                f"/h100/{case_id}_2gpu_frame_mid.png"
             ),
         ),
         (
-            f"{case_id}_2gpu_frame_last.png",
+            f"h100/{case_id}_2gpu_frame_last.png",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{case_id}_2gpu_frame_last.png"
+                f"/h100/{case_id}_2gpu_frame_last.png"
             ),
         ),
     ]
@@ -178,7 +178,10 @@ def test_ltx_hq_remote_gt_uses_sglang_generated_when_official_declared(monkeypat
     files = test_utils._find_remote_consistency_gt_files(case_id, 1, is_video=True)
 
     assert files == [
-        (filename, f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}/{filename}")
+        (
+            f"h100/{filename}",
+            f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}/h100/{filename}",
+        )
         for filename in filenames
     ]
 
@@ -201,10 +204,10 @@ def test_remote_image_gt_falls_back_to_sglang_when_official_missing(monkeypatch)
 
     assert files == [
         (
-            expected_filename,
+            f"h100/{expected_filename}",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{expected_filename}"
+                f"/h100/{expected_filename}"
             ),
         )
     ]
@@ -230,10 +233,10 @@ def test_remote_image_gt_skips_official_for_quarantined_case(monkeypatch):
 
     assert files == [
         (
-            expected_filename,
+            f"h100/{expected_filename}",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{expected_filename}"
+                f"/h100/{expected_filename}"
             ),
         )
     ]
@@ -306,24 +309,24 @@ def test_remote_video_gt_skips_official_for_quarantined_case(monkeypatch):
 
     assert files == [
         (
-            f"{case_id}_2gpu_frame_0.png",
+            f"h100/{case_id}_2gpu_frame_0.png",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{case_id}_2gpu_frame_0.png"
+                f"/h100/{case_id}_2gpu_frame_0.png"
             ),
         ),
         (
-            f"{case_id}_2gpu_frame_mid.png",
+            f"h100/{case_id}_2gpu_frame_mid.png",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{case_id}_2gpu_frame_mid.png"
+                f"/h100/{case_id}_2gpu_frame_mid.png"
             ),
         ),
         (
-            f"{case_id}_2gpu_frame_last.png",
+            f"h100/{case_id}_2gpu_frame_last.png",
             (
                 f"{test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE}"
-                f"/{case_id}_2gpu_frame_last.png"
+                f"/h100/{case_id}_2gpu_frame_last.png"
             ),
         ),
     ]
@@ -377,6 +380,39 @@ def test_platform_gt_candidates_prefer_platform_then_default(monkeypatch):
     ]
 
 
+def test_h100_gt_candidates_prefer_platform_then_default(monkeypatch):
+    monkeypatch.setenv(test_utils.CONSISTENCY_PLATFORM_ENV, "h100")
+
+    assert test_utils.get_consistency_gt_candidates(
+        "unit_image",
+        1,
+        is_video=False,
+        output_format="png",
+    ) == [
+        "h100/unit_image_1gpu.png",
+        "h100/unit_image_1gpu.jpg",
+        "h100/unit_image_1gpu.webp",
+        "unit_image_1gpu.png",
+        "unit_image_1gpu.jpg",
+        "unit_image_1gpu.webp",
+    ]
+
+
+def test_consistency_gt_case_alias_reuses_canonical_filename(monkeypatch):
+    monkeypatch.setenv(test_utils.CONSISTENCY_PLATFORM_ENV, "h100")
+
+    assert test_utils.get_consistency_gt_candidates(
+        "fsdp-inference",
+        2,
+        is_video=False,
+        output_format=None,
+    )[:3] == [
+        "h100/zimage_image_t2i_2_gpus_2gpu.jpg",
+        "h100/zimage_image_t2i_2_gpus_2gpu.png",
+        "h100/zimage_image_t2i_2_gpus_2gpu.webp",
+    ]
+
+
 def test_threshold_metadata_merges_platform_override():
     metadata = test_utils._merge_threshold_metadata(
         {

python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py

@@ -0,0 +1,50 @@
+from types import SimpleNamespace
+
+import torch
+
+from sglang.multimodal_gen.runtime.models.encoders.qwen3 import Qwen3ForCausalLM
+
+
+class _CaptureLayer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.position_ids = None
+        self.attention_lengths = None
+
+    def forward(self, position_ids, hidden_states, residual, attention_lengths):
+        self.position_ids = position_ids
+        self.attention_lengths = attention_lengths
+        if residual is None:
+            residual = torch.zeros_like(hidden_states)
+        return hidden_states, residual
+
+
+class _IdentityNorm(torch.nn.Module):
+    def forward(self, hidden_states, residual):
+        if residual is not None:
+            hidden_states = hidden_states + residual
+        return hidden_states, None
+
+
+def test_default_position_ids_batch_shape():
+    model = Qwen3ForCausalLM.__new__(Qwen3ForCausalLM)
+    torch.nn.Module.__init__(model)
+    layer = _CaptureLayer()
+    model.config = SimpleNamespace(output_hidden_states=False)
+    model.layers = torch.nn.ModuleList([layer])
+    model.norm = _IdentityNorm()
+
+    def get_input_embeddings(input_ids):
+        return torch.zeros(input_ids.shape[0], input_ids.shape[1], 8)
+
+    model.get_input_embeddings = get_input_embeddings
+
+    input_ids = torch.zeros(2, 4, dtype=torch.long)
+    attention_mask = torch.ones(2, 4, dtype=torch.long)
+
+    model(input_ids=input_ids, attention_mask=attention_mask)
+
+    assert layer.position_ids.shape == input_ids.shape
+    assert torch.equal(layer.position_ids[0], torch.arange(4))
+    assert torch.equal(layer.position_ids[1], torch.arange(4))
+    assert layer.attention_lengths == (4, 4)

python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py

@@ -5,13 +5,73 @@ from unittest.mock import patch
 import torch
 
 from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
+from sglang.multimodal_gen.runtime.models.dits.zimage import (
+    ZImageRMSNorm,
+    ZImageTransformer2DModel,
+)
 
 
 class TestZImagePipelineConfig(unittest.TestCase):
+    def test_rmsnorm_native_formula(self) -> None:
+        norm = ZImageRMSNorm(4, eps=1e-5)
+        with torch.no_grad():
+            norm.weight.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0]))
+        x = torch.tensor(
+            [[1.25, 0.5, -0.75, 3.0], [0.1, 2.3, -4.1, 0.7]],
+            dtype=torch.bfloat16,
+        )
+
+        output = norm(x)
+        expected = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + 1e-5)
+        expected = expected * norm.weight.to(dtype=x.dtype)
+
+        self.assertEqual(output.dtype, x.dtype)
+        self.assertTrue(torch.equal(output, expected))
+
+    def test_explicit_sigmas(self) -> None:
+        """Z-Image uses the native explicit flow sigmas schedule."""
+        config = ZImagePipelineConfig()
+
+        self.assertEqual(
+            config.prepare_sigmas(None, 4).tolist(),
+            [1.0, 0.75, 0.5, 0.25],
+        )
+
+    def test_autocast_disabled(self) -> None:
+        """Official Z-Image runs bf16 weights without an outer autocast context."""
+        self.assertFalse(ZImagePipelineConfig().enable_autocast)
+
     @patch("sglang.multimodal_gen.configs.pipeline_configs.zimage.get_sp_world_size")
-    def test_zimage_negative_prompt_rotary_embeddings_use_negative_prompt_len(
+    def test_image_rope_patch_tokens(self, mock_get_sp_world_size) -> None:
-        self, mock_get_sp_world_size
+        mock_get_sp_world_size.return_value = 1
-    ) -> None:
+
+        config = ZImagePipelineConfig()
+        config.vae_config.post_init()
+        batch = SimpleNamespace(
+            prompt_embeds=[torch.ones(113, 2560)],
+            prompt_seq_lens=[[113]],
+            negative_prompt_embeds=None,
+            height=480,
+            width=640,
+        )
+
+        def rotary_emb(pos_ids):
+            return pos_ids
+
+        _, image_pos_ids = config.prepare_pos_cond_kwargs(
+            batch=batch,
+            device=torch.device("cpu"),
+            rotary_emb=rotary_emb,
+            dtype=torch.float32,
+        )["freqs_cis"]
+
+        self.assertEqual(image_pos_ids.shape, (1216, 3))
+        self.assertEqual(image_pos_ids[0].tolist(), [129, 0, 0])
+        self.assertEqual(image_pos_ids[1199].tolist(), [129, 29, 39])
+        self.assertEqual(image_pos_ids[-1].tolist(), [0, 0, 0])
+
+    @patch("sglang.multimodal_gen.configs.pipeline_configs.zimage.get_sp_world_size")
+    def test_negative_rope_len(self, mock_get_sp_world_size) -> None:
         """Negative CFG branch should build RoPE positions from negative prompt embeds."""
         mock_get_sp_world_size.return_value = 1
 
@@ -42,6 +102,36 @@ class TestZImagePipelineConfig(unittest.TestCase):
         self.assertEqual(cap_pos_ids.shape, (neg_cap_padded_len, 3))
         self.assertEqual(image_pos_ids[0].tolist(), [neg_cap_padded_len + 1, 0, 0])
 
+    def test_batched_rope_offsets(self) -> None:
+        model = ZImageTransformer2DModel.__new__(ZImageTransformer2DModel)
+
+        def rotary_emb(pos_ids):
+            return (
+                pos_ids.to(torch.float32),
+                (pos_ids + 1000).to(torch.float32),
+            )
+
+        model.rotary_emb = rotary_emb
+
+        images = [torch.zeros(16, 1, 60, 80), torch.zeros(16, 1, 60, 80)]
+        cap_feats = [torch.zeros(113, 2560), torch.zeros(177, 2560)]
+
+        cap_freqs, image_freqs = model._build_batched_freqs_cis(
+            images,
+            cap_feats,
+            patch_size=2,
+            f_patch_size=1,
+            image_target_len=1216,
+            cap_target_len=192,
+        )
+
+        self.assertEqual(cap_freqs[0].shape, (2, 192, 3))
+        self.assertEqual(image_freqs[0].shape, (2, 1216, 3))
+        self.assertEqual(image_freqs[0][0, 0].tolist(), [129.0, 0.0, 0.0])
+        self.assertEqual(image_freqs[0][1, 0].tolist(), [193.0, 0.0, 0.0])
+        self.assertEqual(cap_freqs[0][0, 127].tolist(), [128.0, 0.0, 0.0])
+        self.assertEqual(cap_freqs[0][0, 128].tolist(), [0.0, 0.0, 0.0])
+
 
 if __name__ == "__main__":
     unittest.main()