diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index cace646d1..87197e5b7 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -689,6 +689,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None): record_shapes=obj.record_shapes, profile_by_stage=obj.profile_by_stage, merge_profiles=obj.merge_profiles, + profile_prefix=obj.profile_prefix, ) 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 cdecebca3..01ffab062 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1332,6 +1332,7 @@ class ProfileReqInput(BaseReq): record_shapes: Optional[bool] = None # Merge profiles from all ranks into a single trace merge_profiles: bool = False + profile_prefix: Optional[str] = None class ProfileReqType(Enum): @@ -1352,6 +1353,7 @@ class ProfileReq(BaseReq): profile_id: Optional[str] = None # Merge profiles from all ranks into a single trace merge_profiles: bool = False + profile_prefix: Optional[str] = None @dataclass diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py index 969158bd9..b230e90f7 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py @@ -54,6 +54,7 @@ class SchedulerProfilerMixin: profile_by_stage: bool, profile_id: str, merge_profiles: bool = False, + profile_prefix: str = "", ) -> ProfileReqOutput: if self.profile_in_progress: return ProfileReqOutput( @@ -74,6 +75,7 @@ class SchedulerProfilerMixin: self.torch_profiler_record_shapes = record_shapes self.profiler_activities = activities self.profile_id = profile_id + self.profile_prefix = profile_prefix if start_step: self.profiler_start_forward_ct = max(start_step, self.forward_ct + 1) @@ -327,6 +329,7 @@ class SchedulerProfilerMixin: recv_req.profile_by_stage, recv_req.profile_id, recv_req.merge_profiles, + recv_req.profile_prefix, ) else: self.init_profile( @@ -339,6 +342,7 @@ class SchedulerProfilerMixin: recv_req.profile_by_stage, recv_req.profile_id, recv_req.merge_profiles, + recv_req.profile_prefix, ) return self.start_profile() else: diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 107a7f6ca..70129ea8c 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -315,6 +315,7 @@ class TokenizerCommunicatorMixin: record_shapes: Optional[bool] = None, profile_by_stage: bool = False, merge_profiles: bool = False, + profile_prefix: Optional[str] = None, ): self.auto_create_handle_loop() env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true") @@ -334,6 +335,7 @@ class TokenizerCommunicatorMixin: profile_by_stage=profile_by_stage, profile_id=str(time.time()), merge_profiles=merge_profiles, + profile_prefix=profile_prefix, ) return await self._execute_profile(req) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a4f2e7025..53969bbf9 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -90,6 +90,27 @@ class ForwardMode(IntEnum): # Split Prefill for PD multiplexing SPLIT_PREFILL = auto() + def __str__(self): + if self == ForwardMode.EXTEND: + return "EXTEND" + elif self == ForwardMode.DECODE: + return "DECODE" + elif self == ForwardMode.MIXED: + return "MIXED" + elif self == ForwardMode.IDLE: + return "IDLE" + elif self == ForwardMode.TARGET_VERIFY: + return "TARGET_VERIFY" + elif self == ForwardMode.DRAFT_EXTEND: + return "DRAFT_EXTEND" + elif self == ForwardMode.DRAFT_EXTEND_V2: + return "DRAFT_EXTEND_V2" + elif self == ForwardMode.PREBUILT: + return "PREBUILT" + elif self == ForwardMode.SPLIT_PREFILL: + return "SPLIT_PREFILL" + return "UNKNOWN" + def is_prefill(self): return self.is_extend() diff --git a/python/sglang/test/send_one.py b/python/sglang/test/send_one.py index cff8ec733..1a9ec1dac 100644 --- a/python/sglang/test/send_one.py +++ b/python/sglang/test/send_one.py @@ -37,6 +37,7 @@ class BenchArgs: profile: bool = False profile_steps: int = 3 profile_by_stage: bool = False + profile_name_prefix: str = "" @staticmethod def add_cli_args(parser: argparse.ArgumentParser): @@ -44,6 +45,9 @@ class BenchArgs: parser.add_argument("--port", type=int, default=BenchArgs.port) parser.add_argument("--batch-size", type=int, default=BenchArgs.batch_size) parser.add_argument("--temperature", type=float, default=BenchArgs.temperature) + parser.add_argument( + "--profile-name-prefix", type=str, default=BenchArgs.profile_name_prefix + ) parser.add_argument( "--max-new-tokens", type=int, default=BenchArgs.max_new_tokens ) @@ -133,6 +137,7 @@ def send_one_prompt(args): None, None, args.profile_by_stage, + args.profile_name_prefix, ) response = requests.post(