diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 0cd227a2c..6346e1583 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/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 ( - 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._maybe_sync_device() 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 ( - 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._maybe_sync_device() execution_time_s = time.perf_counter() - self.start_time if exc_type: diff --git a/test/registered/kernels/ops/diffusion/test_stage_profiler_sync.py b/test/registered/kernels/ops/diffusion/test_stage_profiler_sync.py new file mode 100644 index 000000000..38993033a --- /dev/null +++ b/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"]))