From fa185ed84d43f2015dce7a71ed1ae73038e071e5 Mon Sep 17 00:00:00 2001 From: Chi McIsaac <153383231+qimcis@users.noreply.github.com> Date: Tue, 7 Jul 2026 18:08:32 -0700 Subject: [PATCH] [diffusion] fix: fix z-Image accuracy (#29742) Co-authored-by: Mick --- .../diffusion/triton/zimage_native_norm.py | 182 ++++++ .../configs/pipeline_configs/zimage.py | 4 + .../runtime/layers/attention/__init__.py | 4 + .../runtime/layers/attention/layer.py | 79 +++ .../runtime/models/dits/zimage.py | 573 ++++++++++++++++-- .../runtime/models/encoders/qwen3.py | 8 +- .../test/server/perf_baselines/h100.json | 112 ++-- .../sglang/multimodal_gen/test/test_utils.py | 18 +- .../test/unit/test_consistency_metrics.py | 84 ++- .../test/unit/test_qwen3_encoder.py | 50 ++ .../test/unit/test_zimage_pipeline_config.py | 96 ++- 11 files changed, 1065 insertions(+), 145 deletions(-) create mode 100644 python/sglang/jit_kernel/diffusion/triton/zimage_native_norm.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py diff --git a/python/sglang/jit_kernel/diffusion/triton/zimage_native_norm.py b/python/sglang/jit_kernel/diffusion/triton/zimage_native_norm.py new file mode 100644 index 000000000..82eb4595a --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index 8b4a6a4dc..11b06cb1c 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py b/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py index def73426c..fde7392a1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py +++ b/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", ] diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 24e82377a..635e42390 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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.""" diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index eda6ea3e0..a0f36dcc3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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_k = RMSNorm(self.head_dim, eps=eps) + 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 @@ -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.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.attention_norm1 = ZImageRMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = ZImageRMSNorm(dim, eps=norm_eps) - self.attention_norm2 = RMSNorm(dim, eps=norm_eps) - self.ffn_norm2 = RMSNorm(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( @@ -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 ( - _is_cuda - and attn_out.is_cuda - and attn_out.shape[-1] % 256 == 0 - and attn_out.shape[-1] <= 8192 - and self.attention_norm2.variance_epsilon - == self.ffn_norm1.variance_epsilon - ): - from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import ( - fused_norm_tanh_mul_add_norm_scale, - ) - - 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) + 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 = 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( - target_image_seq_len, + image_attn_len = max( + 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[pad_mask] = pad_token.to(device=tensor.device, dtype=tensor.dtype) + tensor = tensor.clone() + 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[1], cap_freqs_cis[1]], 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=-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, ) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py index 0b19d9f34..8c90cc11d 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py +++ b/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( - 0, hidden_states.shape[1], device=hidden_states.device - ).unsqueeze(0) + position_ids = ( + torch.arange(0, hidden_states.shape[1], device=hidden_states.device) + .unsqueeze(0) + .expand(hidden_states.shape[0], -1) + ) attention_lengths = None if attention_mask is not None: diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 9437078ec..21508bcf9 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/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, - "TextEncodingStage": 130.38, - "LatentPreparationStage": 0.14, - "TimestepPreparationStage": 31.38 + "DenoisingStage": 771.7, + "TextEncodingStage": 130.68, + "LatentPreparationStage": 0.16, + "TimestepPreparationStage": 27.67 }, "denoise_step_ms": { - "0": 19.93, - "1": 28.65, - "2": 83.91, - "3": 83.48, - "4": 83.51, - "5": 83.69, - "6": 84.08, - "7": 84.04, - "8": 84.35 + "0": 95.97, + "1": 29.8, + "2": 91.06, + "3": 91.41, + "4": 93.26, + "5": 90.11, + "6": 90.88, + "7": 91.93, + "8": 93.16 }, - "expected_e2e_ms": 1027.94, - "expected_avg_denoise_ms": 74.6, - "expected_median_denoise_ms": 86.63, - "estimated_full_test_time_s": 121.1 + "expected_e2e_ms": 1068.66, + "expected_avg_denoise_ms": 85.29, + "expected_median_denoise_ms": 91.41, + "estimated_full_test_time_s": 116.3 }, "zimage_image_t2i_fp8": { "stages_ms": { - "TextEncodingStage": 129.84, - "DenoisingStage": 634.42, - "InputValidationStage": 0.03, - "LatentPreparationStage": 0.11, - "TimestepPreparationStage": 17.42, - "DecodingStage": 9.46 + "InputValidationStage": 0.04, + "TextEncodingStage": 131.15, + "LatentPreparationStage": 0.18, + "TimestepPreparationStage": 26.29, + "DenoisingStage": 959.66, + "DecodingStage": 125.62 }, "denoise_step_ms": { - "0": 33.35, - "1": 36.43, - "2": 66.88, - "3": 78.15, - "4": 78.0, - "5": 78.36, - "6": 78.48, - "7": 78.32, - "8": 73.36 + "0": 115.4, + "1": 56.32, + "2": 111.77, + "3": 110.02, + "4": 111.43, + "5": 111.32, + "6": 110.77, + "7": 110.53, + "8": 118.39 }, - "expected_e2e_ms": 958.32, - "expected_avg_denoise_ms": 70.04, - "expected_median_denoise_ms": 81.35, - "estimated_full_test_time_s": 121.0 + "expected_e2e_ms": 1247.73, + "expected_avg_denoise_ms": 106.22, + "expected_median_denoise_ms": 111.32, + "estimated_full_test_time_s": 123.7 }, "zimage_image_t2i_multi_lora": { "stages_ms": { - "TimestepPreparationStage": 24.93, - "DenoisingStage": 673.95, - "DecodingStage": 8.43, - "LatentPreparationStage": 0.1, - "TextEncodingStage": 129.4, - "InputValidationStage": 0.03 + "InputValidationStage": 0.04, + "TextEncodingStage": 130.25, + "LatentPreparationStage": 0.12, + "TimestepPreparationStage": 26.5, + "DenoisingStage": 841.93, + "DecodingStage": 142.78 }, "denoise_step_ms": { - "0": 25.08, - "1": 35.42, - "2": 77.56, - "3": 83.39, - "4": 82.18, - "5": 83.1, - "6": 82.39, - "7": 83.34, - "8": 85.82 + "0": 102.18, + "1": 40.06, + "2": 99.39, + "3": 100.33, + "4": 103.95, + "5": 93.48, + "6": 99.6, + "7": 99.38, + "8": 99.77 }, - "expected_e2e_ms": 1047.82, - "expected_avg_denoise_ms": 74.33, - "expected_median_denoise_ms": 86.33, - "estimated_full_test_time_s": 121.1 + "expected_e2e_ms": 1145.53, + "expected_avg_denoise_ms": 93.12, + "expected_median_denoise_ms": 99.6, + "estimated_full_test_time_s": 162.1 }, "zimage_image_t2i_2_gpus": { "stages_ms": { diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 5391d94d7..f334102d8 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/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 platform == "h100": + if _is_ascend_consistency_case(case_id) or current_platform.is_npu(): 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, ...]: diff --git a/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py b/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py index 34dfe9447..61b79e9fd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py +++ b/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", - "unit_video_1gpu_frame_mid.png", - "unit_video_1gpu_frame_last.png", + "h100/unit_video_1gpu_frame_0.png", + "h100/unit_video_1gpu_frame_mid.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( { diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py b/python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py new file mode 100644 index 000000000..d7e159338 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py index 0a5ed7deb..49fa50db6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py +++ b/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( - self, mock_get_sp_world_size - ) -> None: + def test_image_rope_patch_tokens(self, mock_get_sp_world_size) -> None: + mock_get_sp_world_size.return_value = 1 + + 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()