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

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)