[diffusion] perf_logger: SYNC_STAGE_PROFILING must drain the GPU queue for stage records too (fixes 2-3x inflated DecodingStage readings) (#34124)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -212,6 +212,23 @@ class StageProfiler:
     def _should_record_as_step(self) -> bool:
         return self.record_as_step or self.stage_name.startswith("denoising_step_")
 
+    def _maybe_sync_device(self):
+        """Drain the device queue when SGLANG_DIFFUSION_SYNC_STAGE_PROFILING=1.
+
+        Called at BOTH the timing start and end, for stage records as well as
+        step records. Historically only step records synced, so a stage that
+        merely launches kernels (e.g. DenoisingStage's tail) leaked its queued
+        GPU work into whichever later stage blocked first — DecodingStage
+        readings came out 2-3x too high. The entry sync attributes queued work
+        to the stage that launched it; the exit sync includes this stage's own
+        queued work. Opt-in diagnostics only: the flag defaults off.
+        """
+        if (
+            os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
+            and torch.get_device_module().is_available()
+        ):
+            torch.get_device_module().synchronize()
+
     def __enter__(self):
         if self.log_stage_start_end:
             msg = f"[{self.stage_name}] started..."
@@ -226,12 +243,7 @@ class StageProfiler:
             self.logger.info(msg)
 
         if (self.log_timing and self.metrics) or self.log_stage_start_end:
-            if (
+            self._maybe_sync_device()
-                os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
-                and self._should_record_as_step()
-                and torch.get_device_module().is_available()
-            ):
-                torch.get_device_module().synchronize()
             self.start_time = time.perf_counter()
 
         return self
@@ -240,12 +252,7 @@ class StageProfiler:
         if not ((self.log_timing and self.metrics) or self.log_stage_start_end):
             return False
 
-        if (
+        self._maybe_sync_device()
-            os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
-            and self._should_record_as_step()
-            and torch.get_device_module().is_available()
-        ):
-            torch.get_device_module().synchronize()
         execution_time_s = time.perf_counter() - self.start_time
 
         if exc_type:

test/registered/kernels/ops/diffusion/test_stage_profiler_sync.py

@@ -0,0 +1,50 @@
+"""SGLANG_DIFFUSION_SYNC_STAGE_PROFILING must drain the GPU queue at the
+timing start of *stage* records too — otherwise a stage that only launches
+kernels (DenoisingStage's tail) leaks its queued work into whichever later
+stage blocks first, inflating e.g. DecodingStage readings 2-3x."""
+
+import sys
+import time
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.perf_logger import (
+    RequestMetrics,
+    StageProfiler,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_stage_entry_sync_excludes_previous_stage_tail(monkeypatch):
+    monkeypatch.setenv("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "1")
+    logger = init_logger(__name__)
+    metrics = RequestMetrics("stage-sync-test")
+
+    # Calibrate ~0.5 s of queued GPU work.
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    torch.cuda._sleep(10_000_000)
+    torch.cuda.synchronize()
+    cycles = int(10_000_000 / max(time.perf_counter() - t0, 1e-9) * 0.5)
+
+    # Producer stage queues work without awaiting it (a denoise tail).
+    with StageProfiler("producer", logger, metrics, perf_dump_path_provided=True):
+        torch.cuda._sleep(cycles)
+    # Consumer stage's first blocking op used to absorb the producer's tail.
+    with StageProfiler("consumer", logger, metrics, perf_dump_path_provided=True):
+        torch.ones(8, device="cuda").sum().cpu()
+
+    producer_ms, consumer_ms = metrics.stages["producer"], metrics.stages["consumer"]
+    assert (
+        producer_ms > 250
+    ), f"queued work not attributed to producer: {metrics.stages}"
+    assert consumer_ms < 100, f"producer tail leaked into consumer: {metrics.stages}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))