diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 4b6da3235..842b8aa1b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/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 = ( - (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) - torch.cuda.current_stream(video.device).synchronize() - del frames - _sendfile_all( - process.stdin.fileno(), - buffer.fd, - (end - start) * height * width * 3, - ) - process.stdin.close() - process.stdin = None - returncode = process.wait() + 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) + ) + frames = frames.permute(1, 2, 3, 0).contiguous() + buffer.tensor[: end - start].copy_( + frames, non_blocking=True + ) + torch.cuda.current_stream(video.device).synchronize() + del frames + _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() diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 2139b876f..1ab442ee4 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index c85af07e9..e1a31976f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/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 - # leaving it in OutputBatch to be pickled later - if is_local_endpoint(self.server_args.scheduler_endpoint): + 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" + ): + 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( - output_batch.output - ) + payload = pickle.dumps(output_batch) - with self._record_return_stage( - output_batch, "Scheduler.return_result.pickle" - ): - payload = pickle.dumps(output_batch) + with self._record_return_stage( + output_batch, "Scheduler.return_result.send" + ): + self.receiver.send_multipart([identity, b"", payload]) - with self._record_return_stage( - output_batch, "Scheduler.return_result.send" - ): - 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, @@ -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}", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py index bcd324efc..3bfd147cf 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py +++ b/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): - payload = self.run_stage_with_context( - stage, payload, server_args, run_stage - ) + 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 + ) return payload @staticmethod diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py index 75b3cc138..4595a2102 100644 --- a/python/sglang/multimodal_gen/runtime/utils/profiler.py +++ b/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