Support profiling only prefill or decode without the other (#14182)

python/sglang/srt/entrypoints/http_server.py

@@ -693,6 +693,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None):
         profile_by_stage=obj.profile_by_stage,
         merge_profiles=obj.merge_profiles,
         profile_prefix=obj.profile_prefix,
+        profile_stages=obj.profile_stages,
     )
     return Response(
         content="Start profiling.\n",

python/sglang/srt/managers/io_struct.py

@@ -1389,6 +1389,8 @@ class ProfileReqInput(BaseReq):
     merge_profiles: bool = False
     # The prefix of the profile filenames
     profile_prefix: Optional[str] = None
+    # Only profile these stages and ignore others
+    profile_stages: Optional[List[str]] = None
 
 
 class ProfileReqType(Enum):
@@ -1409,6 +1411,7 @@ class ProfileReq(BaseReq):
     profile_id: Optional[str] = None
     merge_profiles: bool = False
     profile_prefix: Optional[str] = None
+    profile_stages: Optional[List[str]] = None
 
 
 @dataclass

python/sglang/srt/managers/scheduler_profiler_mixin.py

@@ -69,6 +69,7 @@ class SchedulerProfilerMixin:
         profile_id: str,
         merge_profiles: bool = False,
         profile_prefix: str = "",
+        profile_stages: Optional[List[str]] = None,
     ) -> ProfileReqOutput:
         if envs.SGLANG_PROFILE_V2.get():
             return self._profile_manager.configure(
@@ -82,6 +83,7 @@ class SchedulerProfilerMixin:
                 profile_id=profile_id,
                 merge_profiles=merge_profiles,
                 profile_prefix=profile_prefix,
+                profile_stages=profile_stages,
             )
 
         if self.profile_in_progress:
@@ -377,6 +379,7 @@ class SchedulerProfilerMixin:
                     recv_req.profile_id,
                     recv_req.merge_profiles,
                     recv_req.profile_prefix,
+                    recv_req.profile_stages,
                 )
             else:
                 self.init_profile(

python/sglang/srt/managers/tokenizer_communicator_mixin.py

@@ -325,6 +325,7 @@ class TokenizerCommunicatorMixin:
         profile_by_stage: bool = False,
         merge_profiles: bool = False,
         profile_prefix: Optional[str] = None,
+        profile_stages: Optional[List[str]] = None,
     ):
         self.auto_create_handle_loop()
         env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true")
@@ -345,6 +346,7 @@ class TokenizerCommunicatorMixin:
             profile_id=str(time.time()),
             merge_profiles=merge_profiles,
             profile_prefix=profile_prefix,
+            profile_stages=profile_stages,
         )
         return await self._execute_profile(req)
 

python/sglang/srt/utils/profile_utils.py

@@ -57,6 +57,7 @@ class ProfileManager:
         profile_id: str,
         merge_profiles: bool,
         profile_prefix: str,
+        profile_stages: Optional[List[str]] = None,
     ):
         # not supported yet
         assert start_step is None
@@ -81,7 +82,7 @@ class ProfileManager:
 
         self.stage_based_trigger.configure(
             num_steps=num_steps,
-            interesting_stages=["prefill", "decode"],
+            interesting_stages=profile_stages or ["prefill", "decode"],
         )
 
         return ProfileReqOutput(success=True, message="Succeeded")