[diffusion] optimize: optimize cosmos3 (#27041)

python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py

@@ -22,7 +22,7 @@ from sglang.multimodal_gen.runtime.distributed import (
 )
 from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
 from sglang.multimodal_gen.runtime.layers.linear import (
     MergedColumnParallelLinear,
     ReplicatedLinear,
@@ -132,13 +132,6 @@ def compute_mrope_position_ids_vision(
 # -----------------------------------------------------------------------------
 
 
-def qwen3_rotate_half(x: torch.Tensor) -> torch.Tensor:
-    """Qwen3/Llama-style rotate_half: split first/second half of head_dim."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
 def qwen3_apply_rotary_pos_emb(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -153,8 +146,19 @@ def qwen3_apply_rotary_pos_emb(
         cos: [1, S, 1, D] or broadcastable
         sin: [1, S, 1, D] or broadcastable
     """
-    q_embed = (q * cos) + (qwen3_rotate_half(q) * sin)
+    half = q.shape[-1] // 2
-    k_embed = (k * cos) + (qwen3_rotate_half(k) * sin)
+    q1 = q[..., :half]
+    q2 = q[..., half:]
+    q_embed = torch.empty_like(q)
+    q_embed[..., :half] = q1 * cos[..., :half] - q2 * sin[..., :half]
+    q_embed[..., half:] = q2 * cos[..., half:] + q1 * sin[..., half:]
+
+    half = k.shape[-1] // 2
+    k1 = k[..., :half]
+    k2 = k[..., half:]
+    k_embed = torch.empty_like(k)
+    k_embed[..., :half] = k1 * cos[..., :half] - k2 * sin[..., :half]
+    k_embed[..., half:] = k2 * cos[..., half:] + k1 * sin[..., half:]
     return q_embed, k_embed
 
 
@@ -555,11 +559,8 @@ class Cosmos3CrossAttention(nn.Module):
         ]
         v = qkv[:, :, self.num_attention_heads + self.num_key_value_heads :, :]
 
-        q = F.rms_norm(
+        q, k = apply_qk_norm(
-            q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
+            q.contiguous(), k.contiguous(), self.norm_q, self.norm_k, self.head_dim
-        )
-        k = F.rms_norm(
-            k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
         )
         q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
 

python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py

@@ -986,14 +986,12 @@ class AutoencoderKLWan(ParallelTiledVAE):
             with forward_context(
                 feat_cache_arg=self._feat_map, feat_idx_arg=self._conv_idx
             ):
+                out_chunks = []
                 for i in range(iter_):
                     feat_idx.set(0)
-                    if i == 0:
+                    first_chunk.set(i == 0)
-                        first_chunk.set(True)
+                    out_chunks.append(self.decoder(x[:, :, i : i + 1, :, :]))
-                    else:
+                out = torch.cat(out_chunks, 2) if len(out_chunks) > 1 else out_chunks[0]
-                        first_chunk.set(False)
-                    outs.append(self.decoder(x[:, :, i : i + 1, :, :]))
-            out = torch.cat(outs, 2)
 
             if self.config.patch_size is not None:
                 out = unpatchify(out, patch_size=self.config.patch_size)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py

@@ -9,14 +9,16 @@ composed to create complete diffusion pipelines.
 """
 
 from abc import ABC, abstractmethod
-from collections.abc import Iterator
+from collections.abc import Iterable, Iterator
 from contextlib import contextmanager
 from dataclasses import replace
 from enum import Enum, auto
 
 import torch
+from tqdm.auto import tqdm
 
 from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_world_rank
 from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
     ComponentUse,
 )
@@ -87,6 +89,21 @@ class PipelineStage(StageDedupMixin, ABC):
         """Logs a debug message with the stage name as a prefix."""
         logger.debug(f"[{self.__class__.__name__}] {msg}", *args)
 
+    def progress_bar(
+        self,
+        iterable: Iterable | None = None,
+        total: int | None = None,
+        *,
+        disable: bool = False,
+        **kwargs,
+    ) -> tqdm:
+        return tqdm(
+            iterable=iterable,
+            total=total,
+            disable=disable or get_world_rank() != 0,
+            **kwargs,
+        )
+
     def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
         """
         Verify the input for the stage.

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -10,14 +10,13 @@ import math
 import os
 import time
 import weakref
-from collections.abc import Callable, Iterable
+from collections.abc import Callable
 from dataclasses import dataclass, field, fields
 from functools import lru_cache
 from typing import Any
 
 import torch
 import torch.nn as nn
-from tqdm.auto import tqdm
 
 from sglang.jit_kernel.nvfp4 import prewarm_nvfp4_jit_modules
 from sglang.multimodal_gen import envs
@@ -1477,16 +1476,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             return kwargs
         return {k: v for k, v in kwargs.items() if k in param_names}
 
-    def progress_bar(
-        self, iterable: Iterable | None = None, total: int | None = None
-    ) -> tqdm:
-        """
-        Create a progress bar for the denoising process.
-        """
-        local_rank = get_world_group().local_rank
-        disable = local_rank != 0
-        return tqdm(iterable=iterable, total=total, disable=disable)
-
     def _predict_noise_with_cfg(
         self,
         current_model: nn.Module,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py

@@ -14,7 +14,6 @@ import numpy as np
 import PIL.Image
 import torch
 import torch.nn as nn
-from tqdm.auto import tqdm
 
 from sglang.multimodal_gen.configs.sample.sampling_params import DataType
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
@@ -610,7 +609,7 @@ class Cosmos3DenoisingStage(PipelineStage):
             f"CFG_parallel={enable_cfg_parallel}, cfg_rank={cfg_rank}"
         )
 
-        progress_bar = tqdm(
+        progress_bar = self.progress_bar(
             enumerate(timesteps),
             total=len(timesteps),
             desc="Denoising",
@@ -882,7 +881,7 @@ class Cosmos3DecodingStage(PipelineStage):
 
     @staticmethod
     def _postprocess_tensor(decoded: torch.Tensor) -> torch.Tensor:
-        return (decoded * 0.5 + 0.5).clamp(0, 1).float()
+        return decoded.mul_(0.5).add_(0.5).clamp_(0, 1).float()
 
     @staticmethod
     def _postprocess_video_np(video: torch.Tensor, is_image_gen: bool) -> np.ndarray:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py

@@ -14,17 +14,14 @@ from __future__ import annotations
 import functools
 import inspect
 import os
-from collections.abc import Iterable
 
 import torch
 import torch.nn as nn
 from diffusers.utils.torch_utils import randn_tensor
-from tqdm.auto import tqdm
 
 from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.distributed import (
     get_local_torch_device,
-    get_world_group,
 )
 from sglang.multimodal_gen.runtime.distributed.communication_op import (
     cfg_model_parallel_all_reduce,
@@ -319,16 +316,6 @@ class MOVADenoisingStage(PipelineStage):
         result.add_check("audio_latents", batch.audio_latents, V.is_tensor)
         return result
 
-    def progress_bar(
-        self, iterable: Iterable | None = None, total: int | None = None
-    ) -> tqdm:
-        """
-        Create a progress bar for the denoising process.
-        """
-        local_rank = get_world_group().local_rank
-        disable = local_rank != 0
-        return tqdm(iterable=iterable, total=total, disable=disable)
-
     def step_profile(self):
         profiler = SGLDiffusionProfiler.get_instance()
         if profiler: