diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 69d0e9926..05a2fd92d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/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) - 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) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index 23cbeede7..2fe771a68 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/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(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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index c004ad728..97b2a7e03 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/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. 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 4b533ed90..6ae58c861 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 9afd4d5a5..09f6d8685 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py index 06b0ab4ab..aa89a7a9d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py +++ b/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: