minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Chi McIsaac
2026-07-08 09:08:32 +08:00
committed by GitHub
co-authored by Mick
parent d7dcdf3efd
commit fa185ed84d
11 changed files with 1065 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/diffusion/triton/zimage_native_norm.py
+182
View File
@@ -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
+4
View File
@@ -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
+4
View File
@@ -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
+79
View File
@@ -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
+518 -55
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py
+5 -3
View File
@@ -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:
python/sglang/multimodal_gen/test/server/perf_baselines/h100.json
+56 -56
View File
@@ -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": {
python/sglang/multimodal_gen/test/test_utils.py
+14 -4
View File
@@ -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, ...]:
python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py
+60 -24
View File
@@ -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(
{
python/sglang/multimodal_gen/test/unit/test_qwen3_encoder.py
+50
View File
@@ -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
+93 -3
View File
@@ -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()