[diffusion] optimize: optimize cosmos3 (#27041)

2026-06-02 19:47:42 +08:00
parent c2eea4d7b3
commit ce7da7397a
6 changed files with 41 additions and 50 deletions
@@ -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)
-    k_embed = (k * cos) + (qwen3_rotate_half(k) * sin)
+    half = q.shape[-1] // 2
+    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, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
-        )
-        k = F.rms_norm(
-            k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
+        q, k = apply_qk_norm(
+            q.contiguous(), k.contiguous(), self.norm_q, self.norm_k, self.head_dim
        )
        q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)

@@ -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(True)
-                    else:
-                        first_chunk.set(False)
-                    outs.append(self.decoder(x[:, :, i : i + 1, :, :]))
-            out = torch.cat(outs, 2)
+                    first_chunk.set(i == 0)
+                    out_chunks.append(self.decoder(x[:, :, i : i + 1, :, :]))
+                out = torch.cat(out_chunks, 2) if len(out_chunks) > 1 else out_chunks[0]

            if self.config.patch_size is not None:
                out = unpatchify(out, patch_size=self.config.patch_size)
@@ -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.
@@ -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,
@@ -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:
@@ -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: