[diffusion] Fix Helios denoising profiler stepping (#34826)

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

@@ -29,6 +29,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
 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.perf_logger import StageProfiler
+from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
@@ -155,6 +156,7 @@ class HeliosChunkedDenoisingStage(PipelineStage):
         """Denoise a single chunk with full timestep loop."""
         batch_size = latents.shape[0]
         do_cfg = guidance_scale > 1.0
+        profiler = SGLDiffusionProfiler.get_instance()
 
         for i, t in enumerate(timesteps):
             with StageProfiler(
@@ -252,6 +254,8 @@ class HeliosChunkedDenoisingStage(PipelineStage):
                         )
 
                 latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                if profiler:
+                    profiler.step_denoising_step()
 
         return latents
 
@@ -286,6 +290,7 @@ class HeliosChunkedDenoisingStage(PipelineStage):
         """Denoise a single chunk using pyramid super-resolution (Stage 2)."""
         batch_size, num_channel, num_frames, height, width = latents.shape
         patch_size = self.transformer.patch_size
+        profiler = SGLDiffusionProfiler.get_instance()
 
         # Downsample to lowest pyramid level
         latents = latents.permute(0, 2, 1, 3, 4).reshape(
@@ -467,6 +472,8 @@ class HeliosChunkedDenoisingStage(PipelineStage):
                         dmd_timesteps=scheduler.timesteps,
                         all_timesteps=timesteps,
                     )[0]
+                    if profiler:
+                        profiler.step_denoising_step()
 
                 step_counter += 1
 

python/sglang/multimodal_gen/test/unit/test_helios_denoising_profiler.py

@@ -0,0 +1,67 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_denoising import (
+    HeliosChunkedDenoisingStage,
+)
+
+
+class _Transformer:
+    def __call__(self, **kwargs):
+        return torch.zeros_like(kwargs["hidden_states"])
+
+
+class _Scheduler:
+    def step(self, noise_pred, timestep, latents, return_dict=False):
+        return (latents,)
+
+
+class _Profiler:
+    def __init__(self):
+        self.steps = 0
+
+    def step_denoising_step(self):
+        self.steps += 1
+
+
+class TestHeliosDenoisingProfiler(unittest.TestCase):
+    def test_stage1_advances_profiler_once_per_timestep(self):
+        stage = HeliosChunkedDenoisingStage.__new__(HeliosChunkedDenoisingStage)
+        stage.transformer = _Transformer()
+        stage.scheduler = _Scheduler()
+        profiler = _Profiler()
+        timesteps = torch.tensor([2.0, 1.0])
+
+        with patch(
+            "sglang.multimodal_gen.runtime.pipelines_core.stages."
+            "model_specific_stages.helios_denoising."
+            "SGLDiffusionProfiler.get_instance",
+            return_value=profiler,
+        ):
+            output = stage._denoise_one_chunk(
+                latents=torch.ones(1, 2),
+                prompt_embeds=torch.ones(1, 2),
+                negative_prompt_embeds=torch.ones(1, 2),
+                timesteps=timesteps,
+                guidance_scale=1.0,
+                indices_hidden_states=None,
+                indices_latents_history_short=None,
+                indices_latents_history_mid=None,
+                indices_latents_history_long=None,
+                latents_history_short=None,
+                latents_history_mid=None,
+                latents_history_long=None,
+                target_dtype=torch.float32,
+                device=torch.device("cpu"),
+                batch=None,
+                scheduler=stage.scheduler,
+            )
+
+        torch.testing.assert_close(output, torch.ones(1, 2))
+        self.assertEqual(profiler.steps, len(timesteps))
+
+
+if __name__ == "__main__":
+    unittest.main()