[Diffusion] Enable Cosmos3 denoising profiling (#27439)

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