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

[diffusion] fix: fix accuracy for some image models (#20679)

Browse Source
This commit is contained in:
Mick
2026-03-22 15:11:57 +08:00
committed by GitHub
parent 47d36d73f7
commit f7fc2c8592
20 changed files with 599 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
View File

@@ -242,6 +242,19 @@ class PipelineConfig:
def prepare_sigmas(self, sigmas, num_inference_steps):
return sigmas
def get_classifier_free_guidance_scale(self, batch, guidance_scale: float) -> float:
return guidance_scale
def postprocess_cfg_noise(
self,
batch,
noise_pred: torch.Tensor,
noise_pred_cond: torch.Tensor,
) -> torch.Tensor:
# Model-specific CFG variants can override this hook
# e.g. Qwen-Image's true-CFG norm matching.
return noise_pred
## For ImageVAEEncodingStage
def preprocess_condition_image(
self, image, target_width, target_height, _vae_image_processor
@@ -314,6 +327,9 @@ class PipelineConfig:
return shape
def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype:
return prompt_dtype
def allow_set_num_frames(self):
return False
@@ -348,6 +364,15 @@ class PipelineConfig:
latents = sequence_model_parallel_all_gather(latents, dim=2)
return latents
def gather_noise_pred_for_sp(self, batch, noise_pred):
noise_pred = self.gather_latents_for_sp(noise_pred)
raw_latent_shape = getattr(batch, "raw_latent_shape", None)
if raw_latent_shape is not None and noise_pred.dim() == 3:
orig_s = raw_latent_shape[1]
if noise_pred.shape[1] > orig_s:
noise_pred = noise_pred[:, :orig_s, :]
return noise_pred
def preprocess_vae_image(self, batch, vae_image_processor):
pass

94
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
View File

@@ -87,6 +87,32 @@ def _resolve_qwen_edit_per_prompt_images(prompt_list, image_list):
return [[image] for image in image_list]
def _shard_qwen_edit_img_cache_for_sp(
img_cache: torch.Tensor, noisy_img_seq_len: int, device: torch.device
) -> torch.Tensor:
noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
condition_img_cache = shard_rotary_emb_for_sp(img_cache[noisy_img_seq_len:, :])
return torch.cat([noisy_img_cache, condition_img_cache], dim=0).to(device=device)
def _shard_qwen_edit_freqs_cis_for_sp(freqs_cis, noisy_img_seq_len, device):
if isinstance(freqs_cis[0], torch.Tensor) and freqs_cis[0].dim() == 2:
img_cache, txt_cache = freqs_cis
return (
_shard_qwen_edit_img_cache_for_sp(img_cache, noisy_img_seq_len, device),
txt_cache,
)
(img_cos, img_sin), (txt_cos, txt_sin) = freqs_cis
return (
(
_shard_qwen_edit_img_cache_for_sp(img_cos, noisy_img_seq_len, device),
_shard_qwen_edit_img_cache_for_sp(img_sin, noisy_img_seq_len, device),
),
(txt_cos, txt_sin),
)
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
def _pack_latents(latents, batch_size, num_channels_latents, height, width):
latents = latents.view(
@@ -144,12 +170,40 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
def prepare_sigmas(self, sigmas, num_inference_steps):
return self._prepare_sigmas(sigmas, num_inference_steps)
def get_classifier_free_guidance_scale(self, batch, guidance_scale: float) -> float:
if batch.true_cfg_scale is not None:
return batch.true_cfg_scale
return guidance_scale
def postprocess_cfg_noise(
self,
batch,
noise_pred: torch.Tensor,
noise_pred_cond: torch.Tensor,
) -> torch.Tensor:
# Qwen-Image follows the official diffusers true-CFG behavior:
# after combining cond/uncond with true_cfg_scale, match the per-token norm
# back to the conditional branch.
if (
batch.true_cfg_scale is None
or batch.true_cfg_scale <= 1.0
or not batch.do_classifier_free_guidance
):
return noise_pred
cond_norm = torch.norm(noise_pred_cond, dim=-1, keepdim=True)
noise_norm = torch.norm(noise_pred, dim=-1, keepdim=True).clamp_min(1e-12)
return noise_pred * (cond_norm / noise_norm)
def prepare_image_processor_kwargs(self, batch, neg=False):
prompt = batch.prompt if not neg else batch.negative_prompt
if prompt:
prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
txt = prompt_template_encode.format(batch.prompt)
return dict(text=[txt], padding=True)
prompt_list = _normalize_prompt_list(prompt)
txt = [
prompt_template_encode.format(cur_prompt) for cur_prompt in prompt_list
]
return dict(text=txt, padding=True)
else:
return {}
@@ -310,11 +364,9 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
)
img_cache, txt_cache = freqs_cis
noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
img_cache = torch.cat(
[noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
).to(device=device)
img_cache, txt_cache = _shard_qwen_edit_freqs_cis_for_sp(
freqs_cis, noisy_img_seq_len, device
)
return {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": (img_cache, txt_cache),
@@ -500,33 +552,11 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
)
if isinstance(freqs_cis[0], torch.Tensor) and freqs_cis[0].dim() == 2:
img_cache, txt_cache = freqs_cis
noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
img_cache = torch.cat(
[noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
).to(device=device)
return {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": (img_cache, txt_cache),
"img_shapes": img_shapes,
}
(img_cos, img_sin), (txt_cos, txt_sin) = freqs_cis
noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])
# concat back the img_cos for input image (since it is not sp-shared later)
img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to(
device=device
)
img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to(
device=device
)
return {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
"freqs_cis": _shard_qwen_edit_freqs_cis_for_sp(
freqs_cis, noisy_img_seq_len, device
),
"img_shapes": img_shapes,
}

8
python/sglang/multimodal_gen/configs/pipeline_configs/sana.py
View File

@@ -112,3 +112,11 @@ class SanaPipelineConfig(SpatialImagePipelineConfig):
def post_denoising_loop(self, latents, batch):
return latents
def shard_latents_for_sp(self, batch, latents):
# Sana's DiT uses local attention kernels and does not preserve semantics
# when spatial latents are sequence-sharded.
return latents, False
def gather_latents_for_sp(self, latents):
return latents

87
python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py
View File

@@ -46,6 +46,7 @@ class ZImagePipelineConfig(ImagePipelineConfig):
task_type: ModelTaskType = ModelTaskType.T2I
dit_config: DiTConfig = field(default_factory=ZImageDitConfig)
vae_config: VAEConfig = field(default_factory=FluxVAEConfig)
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (Qwen3TextConfig(),)
)
@@ -62,19 +63,22 @@ class ZImagePipelineConfig(ImagePipelineConfig):
F_PATCH_SIZE: int = 1
def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
# flatten to 1-d list
inputs = tokenizer.apply_chat_template(
prompts,
tokenize=True,
add_generation_prompt=True,
enable_thinking=True,
rendered_prompts = [
tokenizer.apply_chat_template(
prompt,
tokenize=False,
add_generation_prompt=True,
enable_thinking=True,
)
for prompt in prompts
]
return tokenizer(
rendered_prompts,
padding="max_length",
max_length=512, # TODO (yhyang201): set max length according to config
truncation=True,
return_tensors="pt",
return_dict=True,
)
return inputs
@staticmethod
def _ceil_to_multiple(x: int, m: int) -> int:
@@ -83,7 +87,7 @@ class ZImagePipelineConfig(ImagePipelineConfig):
return int(math.ceil(x / m) * m)
def _build_zimage_sp_plan(self, batch) -> dict:
"""Build a minimal SP plan on batch for zimage (spatial sharding + cap sharding)."""
"""Build a minimal SP plan on batch for zimage spatial sharding."""
sp_size = get_sp_world_size()
rank = get_sp_parallel_rank()
@@ -112,16 +116,6 @@ class ZImagePipelineConfig(ImagePipelineConfig):
H_tok_local = H_tok_pad // sp_size
h0_tok = rank * H_tok_local
# Cap/text sharding: avoid duplicating cap tokens across ranks.
cap_len = (
int(batch.prompt_embeds[0].size(0))
if getattr(batch, "prompt_embeds", None)
else 0
)
cap_total = self._ceil_to_multiple(cap_len, self.SEQ_LEN_MULTIPLE * sp_size)
cap_local = cap_total // sp_size
cap_start = rank * cap_local
plan = {
"sp_size": sp_size,
"rank": rank,
@@ -135,9 +129,6 @@ class ZImagePipelineConfig(ImagePipelineConfig):
"H_tok_pad": H_tok_pad,
"H_tok_local": H_tok_local,
"h0_tok": h0_tok,
"cap_total": cap_total,
"cap_local": cap_local,
"cap_start": cap_start,
}
batch._zimage_sp_plan = plan
return plan
@@ -149,25 +140,13 @@ class ZImagePipelineConfig(ImagePipelineConfig):
plan = self._build_zimage_sp_plan(batch)
return plan
def _shard_cap(self, cap: torch.Tensor, plan: dict) -> torch.Tensor:
"""cap: [L, D] -> [cap_local, D], padded by repeating last token."""
if plan["sp_size"] <= 1:
return cap
# print(f"cap shape: {cap.shape}") # [L, 2560] for zimage-turbo
L = cap.size(0)
cap_total = plan["cap_total"]
if cap_total > L:
cap = torch.cat([cap, cap[-1:].repeat(cap_total - L, 1)], dim=0)
start = plan["cap_start"]
local = plan["cap_local"]
return cap[start : start + local]
def get_pos_prompt_embeds(self, batch):
# Keep ZImage model signature: encoder_hidden_states is List[Tensor]
if get_sp_world_size() <= 1:
return batch.prompt_embeds
plan = self._get_zimage_sp_plan(batch)
return [self._shard_cap(batch.prompt_embeds[0], plan)]
return batch.prompt_embeds
def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype:
# Match the official diffusers Z-Image pipeline, which samples latents in fp32
# and keeps scheduler state in fp32.
return torch.float32
def shard_latents_for_sp(self, batch, latents):
sp_size = get_sp_world_size()
@@ -203,6 +182,19 @@ class ZImagePipelineConfig(ImagePipelineConfig):
return latents
return sequence_model_parallel_all_gather(latents, dim=3)
def gather_noise_pred_for_sp(self, batch, noise_pred):
# Z-Image shards 5D latents on the effective-H axis, but ComfyUI noise_pred is 4D [B, C, H_local, W].
noise_pred = self.gather_latents_for_sp(noise_pred)
if noise_pred.dim() == 4:
# reconstruct the full spatial tensor
noise_pred = sequence_model_parallel_all_gather(
noise_pred.contiguous(), dim=2
)
# restore the original H/W orientation
if getattr(batch, "_zimage_sp_swap_hw", False):
noise_pred = noise_pred.transpose(2, 3).contiguous()
return noise_pred
def post_denoising_loop(self, latents, batch):
# Restore swapped H/W and crop padded spatial dims before final reshape.
if latents.dim() == 5 and getattr(batch, "_zimage_sp_swap_hw", False):
@@ -233,24 +225,27 @@ class ZImagePipelineConfig(ImagePipelineConfig):
sp_size = get_sp_world_size()
if sp_size > 1:
# SP path: build local-only freqs_cis matching local cap/x.
# SP path: keep caption replicated on every rank and build local-only
# image freqs_cis matching the spatial shard.
plan = self._get_zimage_sp_plan(batch)
cap_ori_len = prompt_embeds.size(0)
cap_padding_len = (-cap_ori_len) % self.SEQ_LEN_MULTIPLE
# cap (local)
# caption (replicated prefix)
cap_pos_ids = create_coordinate_grid(
size=(plan["cap_local"], 1, 1),
start=(1 + plan["cap_start"], 0, 0),
size=(cap_ori_len + cap_padding_len, 1, 1),
start=(1, 0, 0),
device=device,
).flatten(0, 2)
cap_freqs_cis = rotary_emb(cap_pos_ids)
# image (local, effective H-shard). Use cap_total for a stable offset across ranks/passes.
# image (local, effective H-shard), offset after the full caption.
F_tokens = 1
H_tokens_local = plan["H_tok_local"]
W_tokens = plan["W_tok"]
img_pos_ids = create_coordinate_grid(
size=(F_tokens, H_tokens_local, W_tokens),
start=(plan["cap_total"] + 1, plan["h0_tok"], 0),
start=(cap_ori_len + cap_padding_len + 1, plan["h0_tok"], 0),
device=device,
).flatten(0, 2)
img_pad_len = (-img_pos_ids.shape[0]) % self.SEQ_LEN_MULTIPLE

8
python/sglang/multimodal_gen/configs/sample/sampling_params.py
View File

@@ -649,7 +649,7 @@ class SamplingParams:
add_argument(
"--prompt-path",
type=str,
help="Path to a text file containing the prompt",
help="Path to a text file containing prompts (one per line)",
)
add_argument(
"--output-file-name",
@@ -744,6 +744,12 @@ class SamplingParams:
dest="guidance_scale_2",
help="Secondary guidance scale for dual-guidance models (e.g., Wan low-noise expert)",
)
add_argument(
"--true-cfg-scale",
type=float,
dest="true_cfg_scale",
help="True CFG scale for models that distinguish distilled guidance from standard CFG (e.g., Qwen-Image)",
)
add_argument(
"--guidance-rescale",
type=float,

43
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
View File

@@ -354,6 +354,7 @@ class USPAttention(nn.Module):
k: torch.Tensor,
v: torch.Tensor,
num_replicated_prefix: int = 0,
num_replicated_suffix: int = 0,
) -> torch.Tensor:
"""
Forward pass for USPAttention.
@@ -364,6 +365,9 @@ class USPAttention(nn.Module):
in FLUX joint attention. These tokens are excluded from the
Ulysses all-to-all so they appear exactly once in the gathered
sequence, preserving correct attention weights.
num_replicated_suffix: number of trailing tokens in q/k/v that are
replicated across all SP ranks, e.g. caption tokens appended
after image tokens in Z-Image joint attention.
Note: Replicated tensors are not supported in this implementation.
When skip_sequence_parallel=True (set at construction time), all SP
@@ -378,10 +382,18 @@ class USPAttention(nn.Module):
return out
sp_size = get_ulysses_parallel_world_size()
if num_replicated_prefix > 0 and num_replicated_suffix > 0:
raise ValueError(
"USPAttention does not support replicated prefix and suffix at the same time."
)
if sp_size > 1 and num_replicated_prefix > 0:
return self._forward_with_replicated_prefix(
q, k, v, ctx_attn_metadata, num_replicated_prefix
)
if sp_size > 1 and num_replicated_suffix > 0:
return self._forward_with_replicated_suffix(
q, k, v, ctx_attn_metadata, num_replicated_suffix
)
# Ulysses-style All-to-All for sequence/head sharding
if sp_size > 1:
@@ -468,3 +480,34 @@ class USPAttention(nn.Module):
out_rep = torch.cat(gathered, dim=2)
return torch.cat([out_rep, out_shard], dim=1)
def _forward_with_replicated_suffix(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
ctx_attn_metadata,
num_rep: int,
) -> torch.Tensor:
"""Ulysses attention where the last num_rep tokens are replicated
across SP ranks and should not be duplicated by the all-to-all."""
if num_rep <= 0:
raise ValueError("num_rep must be positive for replicated suffix.")
q_shard, q_rep = q[:, :-num_rep], q[:, -num_rep:]
k_shard, k_rep = k[:, :-num_rep], k[:, -num_rep:]
v_shard, v_rep = v[:, :-num_rep], v[:, -num_rep:]
# dense self-attention is permutation equivariant for non-causal use.
# 1. rotate the replicated suffix to the front
# 2. reuse the validated replicated-prefix path, then
# 3. rotate the output back
out = self._forward_with_replicated_prefix(
torch.cat([q_rep, q_shard], dim=1),
torch.cat([k_rep, k_shard], dim=1),
torch.cat([v_rep, v_shard], dim=1),
ctx_attn_metadata,
num_rep,
)
out_rep, out_shard = out[:, :num_rep], out[:, num_rep:]
return torch.cat([out_shard, out_rep], dim=1)

2
python/sglang/multimodal_gen/runtime/layers/attention/selector.py
View File

@@ -137,6 +137,8 @@ def _cached_get_attn_backend(
if len(supported_attention_backends) == 0:
# all attention backends are allowed
pass
elif selected_backend is None and len(supported_attention_backends) == 1:
selected_backend = next(iter(supported_attention_backends))
elif selected_backend is None:
logger.debug(f"Attention backend not specified")
elif selected_backend not in supported_attention_backends:

3
python/sglang/multimodal_gen/runtime/layers/layernorm.py
View File

@@ -549,6 +549,9 @@ def apply_qk_norm(
_is_cuda
and allow_inplace
and (q_eps == k_eps)
and q.dtype in (torch.float16, torch.bfloat16)
and q_norm.weight.dtype == q.dtype
and k_norm.weight.dtype == k.dtype
and can_use_fused_inplace_qknorm(head_dim, q.dtype)
):
fused_inplace_qknorm(

23
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
View File

@@ -9,6 +9,7 @@ import torch.distributed as dist
import torch.nn as nn
from torch import nn
from torch.distributed import init_device_mesh
from transformers import AutoModel
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
from sglang.multimodal_gen.configs.models import EncoderConfig, ModelConfig
@@ -80,6 +81,28 @@ class TextEncoderLoader(ComponentLoader):
use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
return use_cpu_offload
def load_native(
self,
component_model_path: str,
server_args: ServerArgs,
transformers_or_diffusers: str,
):
if transformers_or_diffusers != "transformers":
return super().load_native(
component_model_path, server_args, transformers_or_diffusers
)
encoder_idx = (
1 if component_model_path.rstrip("/").endswith("text_encoder_2") else 0
)
encoder_dtype = server_args.pipeline_config.text_encoder_precisions[encoder_idx]
return AutoModel.from_pretrained(
component_model_path,
trust_remote_code=server_args.trust_remote_code,
revision=server_args.revision,
torch_dtype=PRECISION_TO_TYPE[encoder_dtype],
)
def _prepare_weights(
self,
model_name_or_path: str,

9
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

@@ -2,6 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
import gc
import logging
import multiprocessing as mp
import os
import time
@@ -198,7 +199,7 @@ class GPUWorker:
pool_overhead_gb = peak_reserved_gb - peak_allocated_gb
logger.info(
logger.debug(
f"Peak GPU memory: {peak_reserved_gb:.2f} GB, "
f"Peak allocated: {peak_allocated_gb:.2f} GB, "
f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb / peak_reserved_gb * 100:.1f}%), "
@@ -249,7 +250,11 @@ class GPUWorker:
"after_forward", peak_snapshot
)
if self.rank == 0 and not req.suppress_logs:
if (
self.rank == 0
and not req.suppress_logs
and logger.isEnabledFor(logging.DEBUG)
):
self.do_mem_analysis(output_batch)
duration_ms = (time.monotonic() - start_time) * 1000

7
python/sglang/multimodal_gen/runtime/models/dits/flux.py
View File

@@ -345,6 +345,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
x: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
freqs_cis=None,
num_replicated_prefix: int = 0,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
query, key, value, encoder_query, encoder_key, encoder_value = (
_get_qkv_projections(self, x, encoder_hidden_states)
@@ -380,6 +381,9 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
query = torch.cat([encoder_query, query], dim=1)
key = torch.cat([encoder_key, key], dim=1)
value = torch.cat([encoder_value, value], dim=1)
num_replicated_prefix = (
num_replicated_prefix or encoder_hidden_states.shape[1]
)
if freqs_cis is not None:
cos, sin = freqs_cis
@@ -394,7 +398,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
query, key, cos_sin_cache, is_neox=False
)
x = self.attn(query, key, value)
x = self.attn(query, key, value, num_replicated_prefix=num_replicated_prefix)
x = x.flatten(2, 3)
x = x.to(query.dtype)
@@ -510,6 +514,7 @@ class FluxSingleTransformerBlock(nn.Module):
residual = hidden_states
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
joint_attention_kwargs = joint_attention_kwargs or {}
joint_attention_kwargs.setdefault("num_replicated_prefix", text_seq_len or 0)
if self.use_nunchaku_structure:
if _nunchaku_fused_ops_available:

12
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
View File

@@ -449,15 +449,9 @@ class GlmImageAttention(torch.nn.Module):
assert (
text_attn_mask.dim() == 2
), "the shape of text_attn_mask should be (batch_size, text_seq_length)"
text_attn_mask = text_attn_mask.float().to(query.device)
mix_attn_mask = torch.ones(
(batch_size, text_seq_length + image_seq_length), device=query.device
)
mix_attn_mask[:, :text_seq_length] = text_attn_mask
mix_attn_mask = mix_attn_mask.unsqueeze(2)
attn_mask_matrix = mix_attn_mask @ mix_attn_mask.transpose(1, 2)
attention_mask = (attn_mask_matrix > 0).unsqueeze(1).to(query.dtype)
hidden_states = self.attn(query, key, value)
hidden_states = self.attn(
query, key, value, num_replicated_prefix=text_seq_length
)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)

29
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

@@ -21,6 +21,9 @@ from sglang.jit_kernel.diffusion.triton.scale_shift import (
)
from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_sp_world_size,
)
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
from sglang.multimodal_gen.runtime.layers.layernorm import (
@@ -56,6 +59,16 @@ except Exception:
NunchakuFeedForward = None
def _local_seq_len(seq_len: int, sp_world_size: int) -> int:
"""get the local seq len, from seq_len padding to the next multiple of sp_world_size, then shard to local"""
if sp_world_size <= 1:
return seq_len
padded_len = seq_len
if padded_len % sp_world_size != 0:
padded_len += sp_world_size - (padded_len % sp_world_size)
return padded_len // sp_world_size
def _get_qkv_projections(
attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
):
@@ -648,6 +661,7 @@ class QwenImageCrossAttention(nn.Module):
joint_query,
joint_key,
joint_value,
num_replicated_prefix=seq_len_txt,
)
# Reshape back
@@ -1088,11 +1102,24 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
@functools.lru_cache(maxsize=50)
def build_modulate_index(self, img_shapes: tuple[int, int, int], device):
sp_world_size = get_sp_world_size()
modulate_index_list = []
for sample in img_shapes:
first_size = sample[0][0] * sample[0][1] * sample[0][2]
total_size = sum(s[0] * s[1] * s[2] for s in sample)
idx = (torch.arange(total_size, device=device) >= first_size).int()
if sp_world_size > 1:
first_local_size = _local_seq_len(first_size)
tail_local_size = _local_seq_len(total_size - first_size)
idx = torch.cat(
[
torch.zeros(first_local_size, device=device, dtype=torch.int),
torch.ones(tail_local_size, device=device, dtype=torch.int),
]
)
else:
idx = (torch.arange(total_size, device=device) >= first_size).int()
modulate_index_list.append(idx)
modulate_index = torch.stack(modulate_index_list)

138
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
View File

@@ -5,9 +5,20 @@ import torch
import torch.nn as nn
from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
from sglang.multimodal_gen.runtime.distributed import (
get_sp_parallel_rank,
get_sp_world_size,
get_tp_world_size,
sequence_model_parallel_all_gather,
)
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_ring_parallel_world_size,
)
from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.attention import (
UlyssesAttention,
USPAttention,
)
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
from sglang.multimodal_gen.runtime.layers.linear import (
ColumnParallelLinear,
@@ -209,11 +220,20 @@ class ZImageAttention(nn.Module):
softmax_scale=None,
causal=False,
)
self.ulysses_attn = UlyssesAttention(
num_heads=self.local_num_heads,
head_size=self.head_dim,
num_kv_heads=self.local_num_kv_heads,
softmax_scale=None,
causal=False,
)
def forward(
self,
hidden_states: torch.Tensor,
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
num_replicated_prefix: int = 0,
num_replicated_suffix: int = 0,
):
if self.use_fused_qkv:
qkv, _ = self.to_qkv(hidden_states)
@@ -263,7 +283,42 @@ class ZImageAttention(nn.Module):
q = _apply_rotary_emb(q, cos, sin, is_neox_style=False)
k = _apply_rotary_emb(k, cos, sin, is_neox_style=False)
hidden_states = self.attn(q, k, v)
if (
num_replicated_suffix > 0
and get_sp_world_size() > 1
and get_ring_parallel_world_size() == 1
):
# the cap (last num_replicated_suffix tokens), as condition, should be replicated
q_shard, q_rep = (
q[:, :-num_replicated_suffix],
q[:, -num_replicated_suffix:],
)
k_shard, k_rep = (
k[:, :-num_replicated_suffix],
k[:, -num_replicated_suffix:],
)
v_shard, v_rep = (
v[:, :-num_replicated_suffix],
v[:, -num_replicated_suffix:],
)
hidden_states, hidden_states_rep = self.ulysses_attn(
q_shard,
k_shard,
v_shard,
replicated_q=q_rep,
replicated_k=k_rep,
replicated_v=v_rep,
)
assert hidden_states_rep is not None
hidden_states = torch.cat([hidden_states, hidden_states_rep], dim=1)
else:
hidden_states = self.attn(
q,
k,
v,
num_replicated_prefix=num_replicated_prefix,
num_replicated_suffix=num_replicated_suffix,
)
hidden_states = hidden_states.flatten(2)
hidden_states, _ = self.to_out[0](hidden_states)
@@ -299,6 +354,10 @@ class ZImageTransformerBlock(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.attention",
)
if not modulation:
# Context refiner runs on fully replicated caption tokens only.
# Bypass Ulysses here to preserve the single-GPU attention semantics.
self.attention.attn.skip_sequence_parallel = True
hidden_dim = int(dim / 3 * 8)
nunchaku_enabled = (
@@ -344,6 +403,8 @@ class ZImageTransformerBlock(nn.Module):
x: torch.Tensor,
freqs_cis: Tuple[torch.Tensor, torch.Tensor],
adaln_input: Optional[torch.Tensor] = None,
num_replicated_prefix: int = 0,
num_replicated_suffix: int = 0,
):
if self.modulation:
assert adaln_input is not None
@@ -358,6 +419,8 @@ class ZImageTransformerBlock(nn.Module):
attn_out = self.attention(
self.attention_norm1(x) * scale_msa,
freqs_cis=freqs_cis,
num_replicated_prefix=num_replicated_prefix,
num_replicated_suffix=num_replicated_suffix,
)
x = x + gate_msa * self.attention_norm2(attn_out)
@@ -369,18 +432,21 @@ class ZImageTransformerBlock(nn.Module):
)
else:
# Attention block
attn_input = self.attention_norm1(x)
attn_out = self.attention(
self.attention_norm1(x),
attn_input,
freqs_cis=freqs_cis,
num_replicated_prefix=num_replicated_prefix,
num_replicated_suffix=num_replicated_suffix,
)
x = x + self.attention_norm2(attn_out)
# FFN block
x = x + self.ffn_norm2(
self.feed_forward(
self.ffn_norm1(x),
)
ffn_input = self.ffn_norm1(x)
ffn_out = self.feed_forward(
ffn_input,
)
x = x + self.ffn_norm2(ffn_out)
return x
@@ -666,10 +732,11 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
pH = pW = patch_size
pF = f_patch_size
device = image.device
all_image_out = []
all_image_size = []
all_cap_feats_out = []
all_image_valid_lens = []
all_cap_valid_lens = []
# ------------ Process Caption ------------
cap_ori_len = cap_feat.size(0)
@@ -681,6 +748,7 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
dim=0,
)
all_cap_feats_out.append(cap_padded_feat)
all_cap_valid_lens.append(cap_ori_len)
# ------------ Process Image ------------
C, F, H, W = image.size()
@@ -701,11 +769,14 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
dim=0,
)
all_image_out.append(image_padded_feat)
all_image_valid_lens.append(image_ori_len)
return (
all_image_out,
all_cap_feats_out,
all_image_size,
all_image_valid_lens,
all_cap_valid_lens,
)
def forward(
@@ -734,41 +805,76 @@ class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
x,
cap_feats,
x_size,
x_valid_lens,
cap_valid_lens,
) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size)
x = torch.cat(x, dim=0)
x, _ = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x)
if x_valid_lens[0] < x.shape[0]:
x[x_valid_lens[0] :] = self.x_pad_token.to(dtype=x.dtype)
x_freqs_cis = freqs_cis[1]
x = x.unsqueeze(0)
x_freqs_cis = x_freqs_cis
for layer in self.noise_refiner:
for layer_id, layer in enumerate(self.noise_refiner):
x = layer(x, x_freqs_cis, adaln_input)
cap_feats = torch.cat(cap_feats, dim=0)
cap_feats, _ = self.cap_embedder(cap_feats)
if cap_valid_lens[0] < cap_feats.shape[0]:
cap_feats[cap_valid_lens[0] :] = self.cap_pad_token.to(
dtype=cap_feats.dtype
)
cap_freqs_cis = freqs_cis[0]
cap_feats = cap_feats.unsqueeze(0)
for layer in self.context_refiner:
cap_feats = layer(cap_feats, cap_freqs_cis)
cap_input_dtype = cap_feats.dtype
for layer_id, layer in enumerate(self.context_refiner):
cap_feats = layer(
cap_feats,
cap_freqs_cis,
)
cap_seq_len = cap_feats.shape[1]
use_full_unified_sequence = (
get_sp_world_size() > 1 and get_ring_parallel_world_size() > 1
)
x_local_seq_len = x.shape[1]
if use_full_unified_sequence:
x = sequence_model_parallel_all_gather(x.contiguous(), dim=1)
x_freqs_cis = (
sequence_model_parallel_all_gather(x_freqs_cis[0].contiguous(), dim=0),
sequence_model_parallel_all_gather(x_freqs_cis[1].contiguous(), dim=0),
)
unified = torch.cat([x, cap_feats], dim=1)
unified_freqs_cis = (
torch.cat([x_freqs_cis[0], cap_freqs_cis[0]], dim=0),
torch.cat([x_freqs_cis[1], cap_freqs_cis[1]], dim=0),
)
num_replicated_suffix = cap_seq_len if not use_full_unified_sequence else 0
for layer in self.layers:
unified = layer(unified, unified_freqs_cis, adaln_input)
for layer_id, layer in enumerate(self.layers):
layer.attention.attn.skip_sequence_parallel = use_full_unified_sequence
unified = layer(
unified,
unified_freqs_cis,
adaln_input,
num_replicated_suffix=num_replicated_suffix,
)
unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](
unified, adaln_input
)
unified = list(unified.unbind(dim=0))
x = self.unpatchify(unified, x_size, patch_size, f_patch_size)
if use_full_unified_sequence:
sp_rank = get_sp_parallel_rank()
start = sp_rank * x_local_seq_len
end = start + x_local_seq_len
unified = unified[:, start:end]
x = list(unified.unbind(dim=0))
x = self.unpatchify(x, x_size, patch_size, f_patch_size)
return -x[0]

12
python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py
View File

@@ -64,6 +64,7 @@ import torch
import torch.nn as nn
from transformers.activations import ACT2FN
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
Qwen2_5_VisionRotaryEmbedding,
Qwen2_5_VisionTransformerPretrainedModel,
Qwen2_5_VLAttention,
Qwen2_5_VLCausalLMOutputWithPast,
@@ -515,6 +516,17 @@ class Qwen2_5_VLModel(nn.Module):
config.vision_config
)
self.visual.to(torch.get_default_dtype())
# keeps the vision rotary frequencies in fp32 even when weights are bf16 (as HF does)
head_dim = (
config.vision_config.hidden_size // config.vision_config.num_heads
)
rotary_dim = head_dim // 2
inv_freq = Qwen2_5_VisionRotaryEmbedding(rotary_dim).inv_freq
self.visual.rotary_pos_emb.register_buffer(
"inv_freq",
inv_freq,
persistent=False,
)
self.rope_deltas = None # cache rope_deltas here
self.config = config
# Initialize weights and apply final processing

9
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
View File

@@ -260,8 +260,13 @@ class Req:
def validate(self):
"""Initialize dependent fields after dataclass initialization."""
# Set do_classifier_free_guidance based on guidance scale and negative prompt
if self.guidance_scale > 1.0 and self.negative_prompt is not None:
# Prefer true_cfg_scale when it is explicitly provided.
cfg_scale = (
self.true_cfg_scale
if self.true_cfg_scale is not None
else self.guidance_scale
)
if cfg_scale > 1.0 and self.negative_prompt is not None:
self.do_classifier_free_guidance = True
if self.negative_prompt_embeds is None:
self.negative_prompt_embeds = []

1
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
View File

@@ -238,6 +238,7 @@ class DecodingStage(PipelineStage):
trajectory_latents=batch.trajectory_latents,
trajectory_decoded=trajectory_decoded,
metrics=batch.metrics,
noise_pred=None,
)
# Keep VAE resident during warmup; the real request needs it next.

254
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

@@ -732,13 +732,9 @@ class DenoisingStage(PipelineStage):
and hasattr(batch, "noise_pred")
and batch.noise_pred is not None
):
batch.noise_pred = server_args.pipeline_config.gather_latents_for_sp(
batch.noise_pred
batch.noise_pred = server_args.pipeline_config.gather_noise_pred_for_sp(
batch, batch.noise_pred
)
if hasattr(batch, "raw_latent_shape"):
orig_s = batch.raw_latent_shape[1]
if batch.noise_pred.shape[1] > orig_s:
batch.noise_pred = batch.noise_pred[:, :orig_s, :]
if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
@@ -1166,7 +1162,7 @@ class DenoisingStage(PipelineStage):
disable = local_rank != 0
return tqdm(iterable=iterable, total=total, disable=disable)
def rescale_noise_cfg(
def _rescale_noise_cfg(
self, noise_cfg, noise_pred_text, guidance_rescale=0.0
) -> torch.Tensor:
"""
@@ -1195,6 +1191,163 @@ class DenoisingStage(PipelineStage):
)
return noise_cfg
def _apply_cfg_normalization(
self,
noise_pred: torch.Tensor,
noise_pred_cond: torch.Tensor,
cfg_normalization: float,
) -> torch.Tensor:
factor = float(cfg_normalization)
cond_f = noise_pred_cond.float()
pred_f = noise_pred.float()
ori_norm = torch.linalg.vector_norm(cond_f)
new_norm = torch.linalg.vector_norm(pred_f)
max_norm = ori_norm * factor
if new_norm > max_norm:
noise_pred = noise_pred * (max_norm / new_norm)
return noise_pred
def _apply_cfg_normalization_parallel(
self,
noise_pred: torch.Tensor,
noise_pred_cond: torch.Tensor | None,
cfg_normalization: float,
cfg_rank: int,
) -> torch.Tensor:
# In cfg-parallel mode, only rank 0 has the conditional branch locally,
# so the reference norm has to be broadcast to the other ranks
factor = float(cfg_normalization)
pred_f = noise_pred.float()
new_norm = torch.linalg.vector_norm(pred_f)
if cfg_rank == 0:
assert noise_pred_cond is not None
ori_norm = torch.linalg.vector_norm(noise_pred_cond.float())
else:
ori_norm = torch.empty_like(new_norm)
ori_norm = get_cfg_group().broadcast(ori_norm, src=0)
max_norm = ori_norm * factor
if new_norm > max_norm:
noise_pred = noise_pred * (max_norm / new_norm)
return noise_pred
def _apply_guidance_rescale_parallel(
self,
noise_pred: torch.Tensor,
noise_pred_cond: torch.Tensor | None,
guidance_rescale: float,
cfg_rank: int,
) -> torch.Tensor:
# Guidance rescale is still defined against the conditional branch, so
# cfg-parallel needs to broadcast that statistic to every rank
std_cfg = noise_pred.std(dim=list(range(1, noise_pred.ndim)), keepdim=True)
if cfg_rank == 0:
assert noise_pred_cond is not None
std_text = noise_pred_cond.std(
dim=list(range(1, noise_pred_cond.ndim)), keepdim=True
)
else:
std_text = torch.empty_like(std_cfg)
std_text = get_cfg_group().broadcast(std_text, src=0)
noise_pred_rescaled = noise_pred * (std_text / std_cfg)
return (
guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_pred
)
def _apply_model_specific_cfg_postprocess(
self,
batch: Req,
noise_pred: torch.Tensor,
noise_pred_cond: torch.Tensor | None,
cfg_rank: int,
) -> torch.Tensor:
# keep model-specific CFG behavior out of the main denoising loop
# for cfg-parallel, broadcast cond noise first so the hook sees the same
# inputs as the serial path.
if cfg_rank == 0:
assert noise_pred_cond is not None
cond_noise = noise_pred_cond
else:
# TODO: cache this?
cond_noise = torch.empty_like(noise_pred)
cond_noise = get_cfg_group().broadcast(cond_noise, src=0)
# qwen-image uses true_cfg_scale, match the per-token norm back to the conditional branch
return self.server_args.pipeline_config.postprocess_cfg_noise(
batch, noise_pred, cond_noise
)
def _combine_cfg_parallel(
self,
batch: Req,
noise_pred_cond: torch.Tensor | None,
noise_pred_uncond: torch.Tensor | None,
cfg_scale: float,
cfg_rank: int,
) -> torch.Tensor:
# cfg-parallel splits cond / uncond across ranks and reconstructs the
# final CFG result with an all-reduce.
if cfg_rank == 0:
assert noise_pred_cond is not None
partial = cfg_scale * noise_pred_cond
else:
assert noise_pred_uncond is not None
partial = (1 - cfg_scale) * noise_pred_uncond
noise_pred = cfg_model_parallel_all_reduce(partial)
if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
noise_pred = self._apply_cfg_normalization_parallel(
noise_pred,
noise_pred_cond,
batch.cfg_normalization,
cfg_rank,
)
if batch.guidance_rescale > 0.0:
noise_pred = self._apply_guidance_rescale_parallel(
noise_pred,
noise_pred_cond,
batch.guidance_rescale,
cfg_rank,
)
return self._apply_model_specific_cfg_postprocess(
batch, noise_pred, noise_pred_cond, cfg_rank
)
def _combine_cfg_serial(
self,
batch: Req,
noise_pred_cond: torch.Tensor,
noise_pred_uncond: torch.Tensor,
cfg_scale: float,
) -> torch.Tensor:
# Serial CFG keeps both branches local and is the reference path that
# model-specific postprocessing hooks should match.
noise_pred = noise_pred_uncond + cfg_scale * (
noise_pred_cond - noise_pred_uncond
)
if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
noise_pred = self._apply_cfg_normalization(
noise_pred,
noise_pred_cond,
batch.cfg_normalization,
)
if batch.guidance_rescale > 0.0:
noise_pred = self._rescale_noise_cfg(
noise_pred,
noise_pred_cond,
guidance_rescale=batch.guidance_rescale,
)
return self.server_args.pipeline_config.postprocess_cfg_noise(
batch, noise_pred, noise_pred_cond
)
def _build_attn_metadata(
self,
i: int,
@@ -1413,7 +1566,6 @@ class DenoisingStage(PipelineStage):
noise_pred_cond, latents
)
if not batch.do_classifier_free_guidance:
# If CFG is disabled, we are done. Return the conditional prediction.
return noise_pred_cond
# negative pass
@@ -1436,80 +1588,26 @@ class DenoisingStage(PipelineStage):
noise_pred_uncond = server_args.pipeline_config.slice_noise_pred(
noise_pred_uncond, latents
)
cfg_scale = server_args.pipeline_config.get_classifier_free_guidance_scale(
batch, current_guidance_scale
)
# Combine predictions
if server_args.enable_cfg_parallel:
# Each rank computes its partial contribution and we sum via all-reduce:
# final = s*cond + (1-s)*uncond
if cfg_rank == 0:
assert noise_pred_cond is not None
partial = current_guidance_scale * noise_pred_cond
else:
assert noise_pred_uncond is not None
partial = (1 - current_guidance_scale) * noise_pred_uncond
noise_pred = cfg_model_parallel_all_reduce(partial)
if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
factor = float(batch.cfg_normalization)
pred_f = noise_pred.float()
new_norm = torch.linalg.vector_norm(pred_f)
if cfg_rank == 0:
cond_f = noise_pred_cond.float()
ori_norm = torch.linalg.vector_norm(cond_f)
else:
ori_norm = torch.empty_like(new_norm)
ori_norm = get_cfg_group().broadcast(ori_norm, src=0)
max_norm = ori_norm * factor
if new_norm > max_norm:
noise_pred = noise_pred * (max_norm / new_norm)
# Guidance rescale: broadcast std(cond) from rank 0, compute std(cfg) locally
if batch.guidance_rescale > 0.0:
std_cfg = noise_pred.std(
dim=list(range(1, noise_pred.ndim)), keepdim=True
)
if cfg_rank == 0:
assert noise_pred_cond is not None
std_text = noise_pred_cond.std(
dim=list(range(1, noise_pred_cond.ndim)), keepdim=True
)
else:
std_text = torch.empty_like(std_cfg)
# Broadcast std_text from local src=0 to all ranks in CFG group
std_text = get_cfg_group().broadcast(std_text, src=0)
noise_pred_rescaled = noise_pred * (std_text / std_cfg)
noise_pred = (
batch.guidance_rescale * noise_pred_rescaled
+ (1 - batch.guidance_rescale) * noise_pred
)
return noise_pred
else:
# Serial CFG: both cond and uncond are available locally
assert noise_pred_cond is not None and noise_pred_uncond is not None
noise_pred = noise_pred_uncond + current_guidance_scale * (
noise_pred_cond - noise_pred_uncond
return self._combine_cfg_parallel(
batch,
noise_pred_cond,
noise_pred_uncond,
cfg_scale,
cfg_rank,
)
if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
factor = float(batch.cfg_normalization)
cond_f = noise_pred_cond.float()
pred_f = noise_pred.float()
ori_norm = torch.linalg.vector_norm(cond_f)
new_norm = torch.linalg.vector_norm(pred_f)
max_norm = ori_norm * factor
if new_norm > max_norm:
noise_pred = noise_pred * (max_norm / new_norm)
if batch.guidance_rescale > 0.0:
noise_pred = self.rescale_noise_cfg(
noise_pred,
noise_pred_cond,
guidance_rescale=batch.guidance_rescale,
)
return noise_pred
assert noise_pred_cond is not None and noise_pred_uncond is not None
return self._combine_cfg_serial(
batch,
noise_pred_cond,
noise_pred_uncond,
cfg_scale,
)
def prepare_sta_param(self, batch: Req, server_args: ServerArgs):
"""

8
python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py
View File

@@ -7,7 +7,9 @@ Latent preparation stage for diffusion pipelines.
from diffusers.utils.torch_utils import randn_tensor
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.distributed import (
get_local_torch_device,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
@@ -55,7 +57,9 @@ class LatentPreparationStage(PipelineStage):
batch_size = batch.batch_size
# Get required parameters
dtype = batch.prompt_embeds[0].dtype
dtype = server_args.pipeline_config.get_latent_dtype(
batch.prompt_embeds[0].dtype
)
device = get_local_torch_device()
generator = batch.generator
latents = batch.latents

25
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py
View File

@@ -57,17 +57,26 @@ def sample_block_noise(
_, ph, pw = patch_size
block_size = ph * pw
# Explicitly use CPU to avoid requiring MAGMA for cholesky on ROCm/CUDA
# Explicitly use CPU to avoid requiring MAGMA on ROCm/CUDA.
#
# For the default Helios stage-2 setting gamma=1/3 with a 2x2 block, the
# covariance has eigenvalues {0, 1+gamma, 1+gamma, 1+gamma} and is therefore
# only positive semidefinite. `MultivariateNormal(covariance_matrix=...)`
# requires a strictly positive-definite matrix and fails in the Cholesky
# factorization path, so sample from the PSD covariance via eigen-decomposition.
cov = (
torch.eye(block_size, device="cpu") * (1 + gamma)
- torch.ones(block_size, block_size, device="cpu") * gamma
)
dist = torch.distributions.MultivariateNormal(
torch.zeros(block_size, device="cpu"), covariance_matrix=cov
torch.eye(block_size, device="cpu", dtype=torch.float64) * (1 + gamma)
- torch.ones(block_size, block_size, device="cpu", dtype=torch.float64) * gamma
)
block_number = batch_size * channel * num_frames * (height // ph) * (width // pw)
noise = dist.sample((block_number,))
cov = 0.5 * (cov + cov.T)
eigvals, eigvecs = torch.linalg.eigh(cov)
eigvals = eigvals.clamp_min(0.0)
transform = eigvecs @ torch.diag(torch.sqrt(eigvals))
base_noise = torch.randn(
block_number, block_size, device="cpu", dtype=torch.float64
)
noise = (base_noise @ transform.T).to(dtype=torch.float32)
noise = noise.view(
batch_size, channel, num_frames, height // ph, width // pw, ph, pw
)