[diffusion] feat: add maybe_record_function profiler spans for request phases (#35922)

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

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

@@ -44,6 +44,7 @@ from sglang.multimodal_gen.configs.sample.sampling_params import (
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
+from sglang.multimodal_gen.runtime.utils.profiler import maybe_record_function
 from sglang.srt.observability.trace import TraceReqContext
 
 logger = init_logger(__name__)
@@ -610,21 +611,29 @@ def _try_save_cuda_video_direct(
                     assert buffer.tensor is not None
                     for start in range(0, num_frames, chunk_frames):
                         end = min(start + chunk_frames, num_frames)
-                        frames = (
+                        with maybe_record_function(
-                            (video[:, start:end] * 255).clamp_(0, 255).to(torch.uint8)
+                            f"VIDEO_CHUNK frames {start}-{end} convert+pipe_to_x264"
-                        )
+                        ):
-                        frames = frames.permute(1, 2, 3, 0).contiguous()
+                            frames = (
-                        buffer.tensor[: end - start].copy_(frames, non_blocking=True)
+                                (video[:, start:end] * 255)
-                        torch.cuda.current_stream(video.device).synchronize()
+                                .clamp_(0, 255)
-                        del frames
+                                .to(torch.uint8)
-                        _sendfile_all(
+                            )
-                            process.stdin.fileno(),
+                            frames = frames.permute(1, 2, 3, 0).contiguous()
-                            buffer.fd,
+                            buffer.tensor[: end - start].copy_(
-                            (end - start) * height * width * 3,
+                                frames, non_blocking=True
-                        )
+                            )
-                process.stdin.close()
+                            torch.cuda.current_stream(video.device).synchronize()
-                process.stdin = None
+                            del frames
-                returncode = process.wait()
+                            _sendfile_all(
+                                process.stdin.fileno(),
+                                buffer.fd,
+                                (end - start) * height * width * 3,
+                            )
+                with maybe_record_function("FFMPEG_FLUSH stdin_close+wait"):
+                    process.stdin.close()
+                    process.stdin = None
+                    returncode = process.wait()
             finally:
                 if process.stdin is not None:
                     process.stdin.close()

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -77,6 +77,7 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import (
     PerformanceLogger,
     capture_memory_snapshot,
 )
+from sglang.multimodal_gen.runtime.utils.profiler import maybe_record_function
 from sglang.multimodal_gen.runtime.utils.realtime_video import (
     RAW_RGB_CONTENT_TYPE,
     build_raw_rgb_frame_batches,
@@ -550,7 +551,9 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
             for metrics in output_metrics:
                 metrics.total_duration_ms = duration_ms
 
-            self._materialize_output_transport(output_batch, req, save_output_paths)
+            req_label = req.request_id[:8] if req.request_id else "unnamed"
+            with maybe_record_function(f"SAVE_OUTPUTS {req_label}"):
+                self._materialize_output_transport(output_batch, req, save_output_paths)
             self._record_output_peak_memory(output_batch)
 
             collect_perf = (
@@ -572,7 +575,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
                 and output_batch.output is None
                 and not req.return_raw_frames
             ):
-                torch.get_device_module().empty_cache()
+                with maybe_record_function("EMPTY_CACHE"):
+                    torch.get_device_module().empty_cache()
 
             if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING:
                 if not req.is_warmup:

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -64,6 +64,7 @@ from sglang.multimodal_gen.runtime.server_warmup import (
 from sglang.multimodal_gen.runtime.utils.common import get_zmq_socket
 from sglang.multimodal_gen.runtime.utils.distributed import broadcast_pyobj
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.profiler import maybe_record_function
 from sglang.multimodal_gen.runtime.utils.trace_wrapper import DiffStage, trace_slice
 
 logger = init_logger(__name__)
@@ -724,25 +725,38 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
         replies to client, only on rank 0
         """
         if not should_not_return and self.receiver is not None and identity is not None:
-            # if the server is local, use temp file to spill the frame array instead of
+            with maybe_record_function("REPLY spill+pickle+send"):
-            # leaving it in OutputBatch to be pickled later
+                # if the server is local, use temp file to spill the frame array
-            if is_local_endpoint(self.server_args.scheduler_endpoint):
+                # instead of leaving it in OutputBatch to be pickled later
+                if is_local_endpoint(self.server_args.scheduler_endpoint):
+                    with self._record_return_stage(
+                        output_batch, "Scheduler.return_result.spill_arrays"
+                    ):
+                        output_batch.output = spill_large_arrays_to_file_refs(
+                            output_batch.output
+                        )
+
                 with self._record_return_stage(
-                    output_batch, "Scheduler.return_result.spill_arrays"
+                    output_batch, "Scheduler.return_result.pickle"
                 ):
-                    output_batch.output = spill_large_arrays_to_file_refs(
+                    payload = pickle.dumps(output_batch)
-                        output_batch.output
-                    )
 
-            with self._record_return_stage(
+                with self._record_return_stage(
-                output_batch, "Scheduler.return_result.pickle"
+                    output_batch, "Scheduler.return_result.send"
-            ):
+                ):
-                payload = pickle.dumps(output_batch)
+                    self.receiver.send_multipart([identity, b"", payload])
 
-            with self._record_return_stage(
+    @staticmethod
-                output_batch, "Scheduler.return_result.send"
+    def _req_label(items: list) -> str:
-            ):
+        """Short request tag for profiler span names."""
-                self.receiver.send_multipart([identity, b"", payload])
+        req = items[0][1] if items else None
+        if isinstance(req, list) and req:
+            req = req[0]
+        # request_id is Optional; server warmup and bare server-test
+        # requests arrive without one.
+        if isinstance(req, Req) and req.request_id:
+            return req.request_id[:8]
+        return type(req).__name__
 
     def _return_item_result(
         self,
@@ -1231,7 +1245,10 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
                 continue
 
             try:
-                handler_result = self._dispatch_items(items)
+                with maybe_record_function(
+                    f"REQ {self._req_label(items)} dispatch+forward"
+                ):
+                    handler_result = self._dispatch_items(items)
             except Exception as e:
                 logger.error(
                     f"Error executing request in scheduler event loop: {e}",

python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py

@@ -19,7 +19,10 @@ 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.nvtx_pytorch_hooks import maybe_nvtx_range
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
-from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
+from sglang.multimodal_gen.runtime.utils.profiler import (
+    SGLDiffusionProfiler,
+    maybe_record_function,
+)
 
 if TYPE_CHECKING:
     # Only for type checkers; avoids runtime circular import
@@ -112,10 +115,11 @@ class PipelineExecutor(ABC):
     ) -> Any:
         stage_name = stage._component_stage_name()
         self.before_stage(stage, stage_index, payload, server_args)
-        with maybe_nvtx_range(f"stage_{stage_name}", use_nvtx):
+        with maybe_record_function(f"STAGE {stage_name}"):
-            payload = self.run_stage_with_context(
+            with maybe_nvtx_range(f"stage_{stage_name}", use_nvtx):
-                stage, payload, server_args, run_stage
+                payload = self.run_stage_with_context(
-            )
+                    stage, payload, server_args, run_stage
+                )
         return payload
 
     @staticmethod

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

@@ -1,5 +1,7 @@
+import contextlib
 import gzip
 import os
+from collections.abc import Iterator
 
 import torch
 
@@ -19,6 +21,21 @@ if current_platform.is_npu():
 logger = init_logger(__name__)
 
 
+@contextlib.contextmanager
+def maybe_record_function(name: str, enabled: bool = True) -> Iterator[None]:
+    """Named ``torch.profiler`` span, near-free when no profiler is active.
+    ``enabled`` mirrors :func:`maybe_nvtx_range` for per-request gates such
+    as warmup exclusion.
+    """
+    # Gate on the process-wide module flag; the thread-local _profiler_enabled()
+    # is False under profile_all_threads=True and on non-initiating threads.
+    if not enabled or not torch.autograd.profiler._is_profiler_enabled:
+        yield
+        return
+    with torch.profiler.record_function(name):
+        yield
+
+
 def _resolve_profiler_log_dir(log_dir: str | None) -> str:
     if log_dir is not None:
         return log_dir