diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 35df3cc8a..7303f6994 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 991fb6ea6..bf058ad05 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -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, } diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py index 4815ab76a..73f62fa00 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index e2e180a97..3bdbf6059 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index e77fb2f90..7dcf9bf1d 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 4bfe9f2c0..2a695d15a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index a82a4eca8..646cbd429 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index b69040672..66f3f1f61 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 3b85e7b5e..2e304daf0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 466fa74b6..8757da74e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 1f2c4b806..1e371fe75 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index f7f223dec..54b02861d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 7b10905b9..91f963042 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index ae0e421b6..ca191075f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -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] diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py index fb04d9ba5..160838522 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 160c2b5e6..bb30bd10b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -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 = [] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index 36ee3c119..006139487 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -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. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index c130f9260..937446899 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -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): """ diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py index dbae650c7..fb9879a5f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py index 29b7a3b13..123122716 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py @@ -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 )