[Auto Sync] Improve profilers and simplify bench_one_batch_server.py (#13866)

python/sglang/srt/managers/io_struct.py

@@ -1321,17 +1321,23 @@ class SetInternalStateReqOutput(BaseReq):
 class ProfileReqInput(BaseReq):
     # The output directory
     output_dir: Optional[str] = None
+    # Specify the steps to start the profiling
+    start_step: Optional[int] = None
     # If set, it profile as many as this number of steps.
     # If it is set, profiling is automatically stopped after this step, and
     # the caller doesn't need to run stop_profile.
-    start_step: Optional[int] = None
     num_steps: Optional[int] = None
+    # The activities to record. The choices are ["CPU", "GPU", "MEM", "RPD"]
     activities: Optional[List[str]] = None
+    # Whether profile by stages (e.g., prefill and decode) separately
     profile_by_stage: bool = False
+    # Whether to record source information (file and line number) for the ops.
     with_stack: Optional[bool] = None
+    # Whether to save information about operator’s input shapes.
     record_shapes: Optional[bool] = None
     # Merge profiles from all ranks into a single trace
     merge_profiles: bool = False
+    # The prefix of the profile filenames
     profile_prefix: Optional[str] = None
 
 
@@ -1351,7 +1357,6 @@ class ProfileReq(BaseReq):
     with_stack: Optional[bool] = None
     record_shapes: Optional[bool] = None
     profile_id: Optional[str] = None
-    # Merge profiles from all ranks into a single trace
     merge_profiles: bool = False
     profile_prefix: Optional[str] = None
 

python/sglang/srt/managers/scheduler_profiler_mixin.py

@@ -31,18 +31,22 @@ class SchedulerProfilerMixin:
         self.torch_profiler_output_dir: Optional[Path] = None
         self.profiler_activities: Optional[List[str]] = None
         self.profile_id: Optional[str] = None
+
         self.profiler_start_forward_ct: Optional[int] = None
         self.profiler_target_forward_ct: Optional[int] = None
-        self.profiler_target_prefill_ct: Optional[int] = None
-        self.profiler_target_decode_ct: Optional[int] = None
+
         self.profiler_prefill_ct: Optional[int] = None
         self.profiler_decode_ct: Optional[int] = None
+        self.profiler_target_prefill_ct: Optional[int] = None
+        self.profiler_target_decode_ct: Optional[int] = None
+
         self.profile_by_stage: bool = False
-        self.profile_steps: Optional[int] = None
         self.profile_in_progress: bool = False
-        self.rpd_profiler = None
         self.merge_profiles = False
 
+        # For ROCM
+        self.rpd_profiler = None
+
     def init_profile(
         self,
         output_dir: Optional[str],
@@ -81,12 +85,11 @@ class SchedulerProfilerMixin:
             self.profiler_start_forward_ct = max(start_step, self.forward_ct + 1)
 
         if num_steps:
-            self.profile_steps = num_steps
             if self.profile_by_stage:
-                self.profiler_target_prefill_ct = num_steps
-                self.profiler_target_decode_ct = num_steps
                 self.profiler_prefill_ct = 0
                 self.profiler_decode_ct = 0
+                self.profiler_target_prefill_ct = num_steps
+                self.profiler_target_decode_ct = num_steps
             elif start_step:
                 self.profiler_target_forward_ct = (
                     self.profiler_start_forward_ct + num_steps
@@ -119,7 +122,7 @@ class SchedulerProfilerMixin:
             activity_map[a] for a in activities if a in activity_map
         ]
 
-        if "RPD" in activities:
+        if "RPD" in activities:  # for ROCM
             from rpdTracerControl import rpdTracerControl
 
             rpdTracerControl.skipCreate()
@@ -217,6 +220,11 @@ class SchedulerProfilerMixin:
 
         self.torch_profiler_output_dir.mkdir(parents=True, exist_ok=True)
 
+        if self.profile_prefix:
+            stage_prefix = self.profile_prefix + "-"
+        else:
+            stage_prefix = ""
+
         stage_suffix = f"-{stage.name}" if stage else ""
         logger.info("Stop profiling" + stage_suffix + "...")
         if self.torch_profiler is not None:
@@ -233,7 +241,12 @@ class SchedulerProfilerMixin:
                 if getattr(self, "moe_ep_size", 1) > 1:
                     filename_parts.append(f"EP-{getattr(self, 'moe_ep_rank', 0)}")
 
-                filename = "-".join(filename_parts) + stage_suffix + ".trace.json.gz"
+                filename = (
+                    stage_prefix
+                    + "-".join(filename_parts)
+                    + stage_suffix
+                    + ".trace.json.gz"
+                )
 
                 self.torch_profiler.export_chrome_trace(
                     os.path.join(self.torch_profiler_output_dir, filename)

python/sglang/srt/managers/tokenizer_manager.py

@@ -433,19 +433,19 @@ class TokenizerManager(TokenizerCommunicatorMixin):
         self.auto_create_handle_loop()
         obj.normalize_batch_and_arguments()
 
-        external_trace_header = None
-        if request:
-            if "trace_context" in request.headers:
-                trace_set_remote_propagate_context(request.headers["trace_context"])
-            else:
-                external_trace_header = extract_trace_headers(request.headers)
+        if self.enable_trace:
+            external_trace_header = None
+            if request:
+                if "trace_context" in request.headers:
+                    trace_set_remote_propagate_context(request.headers["trace_context"])
+                else:
+                    external_trace_header = extract_trace_headers(request.headers)
+
+            self._trace_request_start(obj, created_time, external_trace_header)
 
         if self.server_args.tokenizer_worker_num > 1:
             self._attach_multi_http_worker_info(obj)
 
-        if self.enable_trace:
-            self._trace_request_start(obj, created_time, external_trace_header)
-
         if self.log_requests:
             max_length, skip_names, _ = self.log_request_metadata
             logger.info(