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

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
 

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)
+            prompt_list = _normalize_prompt_list(prompt)
-            return dict(text=[txt], padding=True)
+            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
+        img_cache, txt_cache = _shard_qwen_edit_freqs_cis_for_sp(
-        noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
+            freqs_cis, noisy_img_seq_len, device
-        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),
@@ -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,
         }
 

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

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
+        rendered_prompts = [
-        inputs = tokenizer.apply_chat_template(
+            tokenizer.apply_chat_template(
-            prompts,
+                prompt,
-            tokenize=True,
+                tokenize=False,
-            add_generation_prompt=True,
+                add_generation_prompt=True,
-            enable_thinking=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]
+        return batch.prompt_embeds
-        if get_sp_world_size() <= 1:
+
-            return batch.prompt_embeds
+    def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype:
-        plan = self._get_zimage_sp_plan(batch)
+        # Match the official diffusers Z-Image pipeline, which samples latents in fp32
-        return [self._shard_cap(batch.prompt_embeds[0], plan)]
+        # 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),
+                size=(cap_ori_len + cap_padding_len, 1, 1),
-                start=(1 + plan["cap_start"], 0, 0),
+                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

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,

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)

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:

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(

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,

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

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:

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)
+        hidden_states = self.attn(
-            mix_attn_mask = torch.ones(
+            query, key, value, num_replicated_prefix=text_seq_length
-                (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 = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.to(query.dtype)
 

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)

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(
+            ffn_input = self.ffn_norm1(x)
-                self.feed_forward(
+            ffn_out = self.feed_forward(
-                    self.ffn_norm1(x),
+                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_input_dtype = cap_feats.dtype
-            cap_feats = layer(cap_feats, cap_freqs_cis)
+        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:
+        for layer_id, layer in enumerate(self.layers):
-            unified = layer(unified, unified_freqs_cis, adaln_input)
+            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))
+        if use_full_unified_sequence:
-        x = self.unpatchify(unified, x_size, patch_size, f_patch_size)
+            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]
 

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

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
+        # Prefer true_cfg_scale when it is explicitly provided.
-        if self.guidance_scale > 1.0 and self.negative_prompt is not None:
+        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 = []

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.

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 = server_args.pipeline_config.gather_noise_pred_for_sp(
-                batch.noise_pred
+                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:
+            return self._combine_cfg_parallel(
-            #   final = s*cond + (1-s)*uncond
+                batch,
-            if cfg_rank == 0:
+                noise_pred_cond,
-                assert noise_pred_cond is not None
+                noise_pred_uncond,
-                partial = current_guidance_scale * noise_pred_cond
+                cfg_scale,
-            else:
+                cfg_rank,
-                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
             )
 
-            if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
+        assert noise_pred_cond is not None and noise_pred_uncond is not None
-                factor = float(batch.cfg_normalization)
+        return self._combine_cfg_serial(
-                cond_f = noise_pred_cond.float()
+            batch,
-                pred_f = noise_pred.float()
+            noise_pred_cond,
-                ori_norm = torch.linalg.vector_norm(cond_f)
+            noise_pred_uncond,
-                new_norm = torch.linalg.vector_norm(pred_f)
+            cfg_scale,
-                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
 
     def prepare_sta_param(self, batch: Req, server_args: ServerArgs):
         """

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

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.eye(block_size, device="cpu", dtype=torch.float64) * (1 + gamma)
-        - torch.ones(block_size, block_size, device="cpu") * gamma
+        - torch.ones(block_size, block_size, device="cpu", dtype=torch.float64) * gamma
-    )
-    dist = torch.distributions.MultivariateNormal(
-        torch.zeros(block_size, device="cpu"), covariance_matrix=cov
     )
     block_number = batch_size * channel * num_frames * (height // ph) * (width // pw)
-
+    cov = 0.5 * (cov + cov.T)
-    noise = dist.sample((block_number,))
+    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
     )