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 84987622a..abeeef8a0 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 @@ -41,6 +41,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler from sglang.srt.utils.common import get_compiler_backend logger = init_logger(__name__) @@ -484,6 +485,11 @@ class Cosmos3DenoisingStage(PipelineStage): result.add_check("timesteps", batch.timesteps, V.is_tensor) return result + def step_profile(self): + profiler = SGLDiffusionProfiler.get_instance() + if profiler: + profiler.step_denoising_step() + def _run_transformer( self, latents: torch.Tensor, @@ -700,6 +706,9 @@ class Cosmos3DenoisingStage(PipelineStage): if image_latent is not None: latents[:, :, 0:1, :, :] = image_latent + if batch.profile and not batch.is_warmup: + self.step_profile() + batch.latents = latents self.log_info("Denoising complete") return batch