[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,11 +611,18 @@ 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)
+                        with maybe_record_function(
+                            f"VIDEO_CHUNK frames {start}-{end} convert+pipe_to_x264"
+                        ):
                             frames = (
-                            (video[:, start:end] * 255).clamp_(0, 255).to(torch.uint8)
+                                (video[:, start:end] * 255)
+                                .clamp_(0, 255)
+                                .to(torch.uint8)
                             )
                             frames = frames.permute(1, 2, 3, 0).contiguous()
-                        buffer.tensor[: end - start].copy_(frames, non_blocking=True)
+                            buffer.tensor[: end - start].copy_(
+                                frames, non_blocking=True
+                            )
                             torch.cuda.current_stream(video.device).synchronize()
                             del frames
                             _sendfile_all(
@@ -622,6 +630,7 @@ def _try_save_cuda_video_direct(
                                 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()

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,6 +551,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
             for metrics in output_metrics:
                 metrics.total_duration_ms = duration_ms
 
+            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)
 
@@ -572,6 +575,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
                 and output_batch.output is None
                 and not req.return_raw_frames
             ):
+                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:

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,8 +725,9 @@ 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
-            # leaving it in OutputBatch to be pickled later
+            with maybe_record_function("REPLY spill+pickle+send"):
+                # if the server is local, use temp file to spill the frame array
+                # 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"
@@ -744,6 +746,18 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
                 ):
                     self.receiver.send_multipart([identity, b"", payload])
 
+    @staticmethod
+    def _req_label(items: list) -> str:
+        """Short request tag for profiler span names."""
+        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,
         item: tuple[bytes | None, Any],
@@ -1231,6 +1245,9 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
                 continue
 
             try:
+                with maybe_record_function(
+                    f"REQ {self._req_label(items)} dispatch+forward"
+                ):
                     handler_result = self._dispatch_items(items)
             except Exception as e:
                 logger.error(

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,6 +115,7 @@ class PipelineExecutor(ABC):
     ) -> Any:
         stage_name = stage._component_stage_name()
         self.before_stage(stage, stage_index, payload, server_args)
+        with maybe_record_function(f"STAGE {stage_name}"):
             with maybe_nvtx_range(f"stage_{stage_name}", use_nvtx):
                 payload = self.run_stage_with_context(
                     stage, payload, server_args, run_stage

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