diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index c94485548..f43843575 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/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", diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 4814786dd..e34736cc4 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py index a929b0ac0..009275f35 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/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( diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index d6e237c71..edbc52526 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/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) diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index 536f46017..c991cba35 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/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")