From 36b729c2b8498c54bfba3fbe579392f9dc9eeb3f Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 30 Nov 2025 16:59:52 +0800 Subject: [PATCH] Implement profiler v2 and fix stage mixture bug (#14148) --- .../srt/managers/scheduler_profiler_mixin.py | 34 ++ python/sglang/srt/utils/profile_utils.py | 380 ++++++++++++++++++ test/srt/run_suite.py | 1 + test/srt/test_profile_v2.py | 102 +++++ 4 files changed, 517 insertions(+) create mode 100644 python/sglang/srt/utils/profile_utils.py create mode 100644 test/srt/test_profile_v2.py diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py index 6cb27441d..a929b0ac0 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py @@ -6,10 +6,12 @@ from typing import List, Optional import torch +from sglang.srt.environ import envs from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileReqType from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.utils import is_npu from sglang.srt.utils.profile_merger import ProfileMerger +from sglang.srt.utils.profile_utils import ProfileManager _is_npu = is_npu() if _is_npu: @@ -27,6 +29,14 @@ logger = logging.getLogger(__name__) class SchedulerProfilerMixin: def init_profiler(self): + if envs.SGLANG_PROFILE_V2.get(): + self._profile_manager = ProfileManager( + tp_rank=self.tp_rank, + cpu_group=self.cpu_group, + gpu_id=self.gpu_id, + ) + return + self.torch_profiler = None self.torch_profiler_output_dir: Optional[Path] = None self.profiler_activities: Optional[List[str]] = None @@ -60,6 +70,20 @@ class SchedulerProfilerMixin: merge_profiles: bool = False, profile_prefix: str = "", ) -> ProfileReqOutput: + if envs.SGLANG_PROFILE_V2.get(): + return self._profile_manager.configure( + output_dir=output_dir, + start_step=start_step, + num_steps=num_steps, + activities=activities, + with_stack=with_stack, + record_shapes=record_shapes, + profile_by_stage=profile_by_stage, + profile_id=profile_id, + merge_profiles=merge_profiles, + profile_prefix=profile_prefix, + ) + if self.profile_in_progress: return ProfileReqOutput( success=False, @@ -105,6 +129,9 @@ class SchedulerProfilerMixin: def start_profile( self, stage: Optional[ForwardMode] = None ) -> ProfileReqOutput | None: + if envs.SGLANG_PROFILE_V2.get(): + return self._profile_manager.manual_start() + stage_str = f" for {stage.name}" if stage else "" logger.info( f"Profiling starts{stage_str}. Traces will be saved to: {self.torch_profiler_output_dir} (with profile id: {self.profile_id})", @@ -212,6 +239,9 @@ class SchedulerProfilerMixin: def stop_profile( self, stage: Optional[ForwardMode] = None ) -> ProfileReqOutput | None: + if envs.SGLANG_PROFILE_V2.get(): + return self._profile_manager.manual_stop() + if not self.profile_in_progress: return ProfileReqOutput( success=False, @@ -294,6 +324,10 @@ class SchedulerProfilerMixin: return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}") def _profile_batch_predicate(self, batch): + if envs.SGLANG_PROFILE_V2.get(): + self._profile_manager.step(forward_mode=batch.forward_mode) + return + if self.profile_by_stage: if batch.forward_mode.is_prefill(): if self.profiler_prefill_ct == 0: diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py new file mode 100644 index 000000000..536f46017 --- /dev/null +++ b/python/sglang/srt/utils/profile_utils.py @@ -0,0 +1,380 @@ +import logging +import os +import time +from abc import ABC +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Dict, List, Optional + +import torch + +from sglang.srt.managers.io_struct import ProfileReqOutput +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.utils import is_npu + +_is_npu = is_npu() +if _is_npu: + import torch_npu + + patches = [ + ["profiler.profile", torch_npu.profiler.profile], + ["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU], + ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU], + ] + torch_npu._apply_patches(patches) + +logger = logging.getLogger(__name__) + + +class ProfileManager: + def __init__(self, tp_rank: int, cpu_group, gpu_id: int): + self.stage_based_trigger = _StageBasedTrigger( + on_start=self._do_start, + on_stop=self._do_stop, + ) + self.tp_rank = tp_rank + self.cpu_group = cpu_group + self.profiler_kwargs = None + self.profiler = None + + def step(self, forward_mode: ForwardMode): + stage = _get_stage_from_forward_mode(forward_mode) + if stage is None: + return + + self.stage_based_trigger.step(stage=stage) + + def configure( + self, + *, + output_dir: Optional[str], + start_step: Optional[int], + num_steps: Optional[int], + activities: Optional[List[str]], + with_stack: Optional[bool], + record_shapes: Optional[bool], + profile_by_stage: bool, + profile_id: str, + merge_profiles: bool, + profile_prefix: str, + ): + # not supported yet + assert start_step is None + assert ( + profile_by_stage + ), "only support profile_by_stage=true now" # `false` can be easily supported + assert not merge_profiles + + if output_dir is None: + output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") + if activities is None: + activities = ["CPU", "GPU"] + + self.profiler_kwargs = dict( + activities=activities, + with_stack=with_stack, + record_shapes=record_shapes, + output_dir=output_dir, + output_prefix=profile_prefix, + profile_id=profile_id, + ) + + self.stage_based_trigger.configure( + num_steps=num_steps, + interesting_stages=["prefill", "decode"], + ) + + return ProfileReqOutput(success=True, message="Succeeded") + + def manual_start(self): + raise NotImplementedError("manually start is only supported yet") + + def manual_stop(self): + raise NotImplementedError("manually stop is only supported yet") + + def _do_start(self, stage: Optional[str] = None): + logger.info( + f"Profiling starts{f' for {stage}' if stage else ''}. " + f"Traces will be saved to: {self.profiler_kwargs['output_dir']} " + f"(with profile id: {self.profiler_kwargs['profile_id']})", + ) + + assert self.profiler is None + self.profiler = _ProfilerBase.create( + **self.profiler_kwargs, + tp_rank=self.tp_rank, + cpu_group=self.cpu_group, + output_suffix=f"-{stage}" if stage else "", + ) + self.profiler.start() + + def _do_stop(self): + logger.info("Stop profiling...") + self.profiler.stop() + logger.info( + f"Profiling done. Traces are saved to: {self.profiler_kwargs['output_dir']}" + ) + self.profiler = None + + +def _get_stage_from_forward_mode(forward_mode: ForwardMode): + if forward_mode.is_prefill(): + return "prefill" + elif forward_mode.is_decode(): + return "decode" + elif forward_mode.is_idle(): + return None + else: + raise RuntimeError(f"unsupported profile stage: {forward_mode=}") + + +# ======================================== Stage related ========================================== + + +class _StageBasedTrigger: + @dataclass + class _StageConfig: + target_count: int + + @dataclass + class _RunningState: + curr_stage: str + curr_count: int + + def __init__(self, on_start: Callable, on_stop: Callable): + self.on_start = on_start + self.on_stop = on_stop + + self.running_state: Optional[_StageBasedTrigger._RunningState] = None + # When a stage is in the dict, it means it is being or should be executed + self.stage_configs: Dict[str, _StageBasedTrigger._StageConfig] = {} + + def configure(self, num_steps: int, interesting_stages: List[str]): + assert self.running_state is None + self.stage_configs = { + stage: self._StageConfig(target_count=num_steps) + for stage in interesting_stages + } + + def step(self, stage: str): + # Incr counter + if (s := self.running_state) is not None: + s.curr_count += 1 + + # Maybe stop + if ((s := self.running_state) is not None) and ( + (s.curr_count > self.stage_configs[s.curr_stage].target_count) + or (stage != s.curr_stage) + ): + del self.stage_configs[s.curr_stage] + self.running_state = None + self.on_stop() + + # Maybe start + if (self.running_state is None) and (stage in self.stage_configs): + self.running_state = self._RunningState( + curr_stage=stage, + curr_count=0, + ) + self.on_start(stage=stage) + + # Sanity check + assert (self.running_state is not None) == (stage in self.stage_configs) + if (s := self.running_state) is not None: + assert s.curr_stage == stage + + +# ======================================== Concrete profilers ========================================== + + +class _ProfilerBase(ABC): + @staticmethod + def create(activities, with_stack, record_shapes, **kwargs): + inners = [] + if ("CPU" in activities) or ("GPU" in activities): + inners.append( + _ProfilerTorch( + **kwargs, + activities=activities, + with_stack=with_stack, + record_shapes=record_shapes, + ) + ) + if "MEM" in activities: + inners.append(_ProfilerMemory(**kwargs)) + if "CUDA_PROFILER" in activities: + inners.append(_ProfilerCudart(**kwargs)) + if "RPD" in activities: # for ROCM + inners.append(_ProfilerRPD(**kwargs)) + + return _ProfilerList(inners) + + def start(self): + raise NotImplementedError + + def stop(self): + raise NotImplementedError + + +class _ProfilerList(_ProfilerBase): + def __init__(self, inners: List[_ProfilerBase]): + self.inners = inners + + def start(self): + for inner in self.inners: + inner.start() + + def stop(self): + for inner in self.inners: + inner.stop() + + +class _ProfilerConcreteBase(_ProfilerBase): + def __init__( + self, + output_dir: str, + output_prefix: str, + output_suffix: str, + profile_id: str, + tp_rank: int, + cpu_group, + ): + self.output_dir = output_dir + self.output_prefix = output_prefix + self.output_suffix = output_suffix + self.profile_id = profile_id + self.tp_rank = tp_rank + self.cpu_group = cpu_group + + +class _ProfilerTorch(_ProfilerConcreteBase): + def __init__(self, with_stack: bool, record_shapes: bool, activities, **kwargs): + super().__init__(**kwargs) + self.with_stack = with_stack + self.record_shapes = record_shapes + self.activities = activities + + def start(self): + activity_map = { + "CPU": torch.profiler.ProfilerActivity.CPU, + "GPU": torch.profiler.ProfilerActivity.CUDA, + } + torchprof_activities = [ + activity_map[a] for a in self.activities if a in activity_map + ] + + self.torch_profiler = torch.profiler.profile( + activities=torchprof_activities, + with_stack=self.with_stack if self.with_stack is not None else True, + record_shapes=( + self.record_shapes if self.record_shapes is not None else False + ), + on_trace_ready=( + None + if not _is_npu + else torch_npu.profiler.tensorboard_trace_handler(self.output_dir) + ), + ) + self.torch_profiler.start() + + def stop(self): + Path(self.output_dir).mkdir(parents=True, exist_ok=True) + + self.torch_profiler.stop() + if not _is_npu: + # Build filename with only non-zero ranks to maintain backward compatibility + filename_parts = [self.profile_id, f"TP-{self.tp_rank}"] + + # Only add other ranks if parallelism is enabled (size > 1) + if getattr(self, "dp_size", 1) > 1: + filename_parts.append(f"DP-{getattr(self, 'dp_rank', 0)}") + if getattr(self, "pp_size", 1) > 1: + filename_parts.append(f"PP-{getattr(self, 'pp_rank', 0)}") + if getattr(self, "moe_ep_size", 1) > 1: + filename_parts.append(f"EP-{getattr(self, 'moe_ep_rank', 0)}") + + filename = ( + (self.output_prefix + "-" if self.output_prefix else "") + + "-".join(filename_parts) + + self.output_suffix + + ".trace.json.gz" + ) + + self.torch_profiler.export_chrome_trace( + os.path.join(self.output_dir, filename) + ) + torch.distributed.barrier(self.cpu_group) + + # TODO: migrate `_merge_profile_traces` + + +class _ProfilerMemory(_ProfilerConcreteBase): + def start(self): + torch.cuda.memory._record_memory_history(max_entries=100000) + + def stop(self): + Path(self.output_dir).mkdir(parents=True, exist_ok=True) + + memory_profile_path = os.path.join( + self.output_dir, + str(time.time()) + + f"-TP-{self.tp_rank}-memory" + + self.output_suffix + + ".pickle", + ) + torch.cuda.memory._dump_snapshot(memory_profile_path) + torch.cuda.memory._record_memory_history(enabled=None) + + +class _ProfilerCudart(_ProfilerConcreteBase): + def start(self): + logger.info(f"Call cudaProfilerStart") + torch.cuda.cudart().cudaProfilerStart() + + def stop(self): + logger.info(f"Call cudaProfilerStop") + torch.cuda.cudart().cudaProfilerStop() + + +class _ProfilerRPD(_ProfilerConcreteBase): + def start(self): + Path(self.output_dir).mkdir(parents=True, exist_ok=True) + + from rpdTracerControl import rpdTracerControl + + rpdTracerControl.skipCreate() + + self.rpd_profile_path = os.path.join( + self.output_dir, + "rpd-" + str(time.time()) + f"-TP-{self.tp_rank}" + ".trace.json.gz", + ) + + if self.tp_rank == 0: + import sqlite3 + + from rocpd.schema import RocpdSchema + + if os.path.exists("trace.rpd"): + os.unlink("trace.rpd") + schema = RocpdSchema() + connection = sqlite3.connect("trace.rpd") + schema.writeSchema(connection) + connection.commit() + del connection + torch.distributed.barrier(self.cpu_group) + + self.rpd_profiler = rpdTracerControl() + self.rpd_profiler.setPythonTrace(True) + self.rpd_profiler.start() + self.rpd_profiler.rangePush("", "rpd profile range", "") + + def stop(self): + self.rpd_profiler.rangePop() + self.rpd_profiler.stop() + self.rpd_profiler.flush() + + torch.distributed.barrier(self.cpu_group) + if self.tp_rank == 0: + from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace + + rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path) diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 13f880d53..e55a1a4dc 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -213,6 +213,7 @@ suites = { TestFile("test_gpt_oss_common.py"), TestFile("test_moe_eval_accuracy_large.py"), TestFile("test_vision_openai_server_common.py"), + TestFile("test_profile_v2.py"), ], } diff --git a/test/srt/test_profile_v2.py b/test/srt/test_profile_v2.py new file mode 100644 index 000000000..8ff16213a --- /dev/null +++ b/test/srt/test_profile_v2.py @@ -0,0 +1,102 @@ +import os +import shutil +import tempfile +import unittest +from pathlib import Path + +import requests + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestStartProfile(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.output_dir = tempfile.mkdtemp() + envs.SGLANG_TORCH_PROFILER_DIR.set(cls.output_dir) + envs.SGLANG_PROFILE_V2.set(True) + cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def setUp(self): + self._clear_profile_dir() + + def test_profile_by_stage(self): + self._start_profile( + profile_by_stage=True, + num_steps=10, + ) + + self._post_request() + + self._check_profile_output(pattern="*-prefill*", expect_existence=True) + self._check_profile_output(pattern="*-decode*", expect_existence=True) + + def test_decode_only(self): + self._start_profile( + profile_by_stage=True, + profile_stages=["decode"], + num_steps=10, + ) + + self._post_request() + + self._check_profile_output(pattern="*-prefill*", expect_existence=False) # NOTE + self._check_profile_output(pattern="*-decode*", expect_existence=True) + + def _start_profile(self, **kwargs): + """Start profiling with optional parameters.""" + response = requests.post( + f"{DEFAULT_URL_FOR_TEST}/start_profile", + json=kwargs if kwargs else None, + ) + self.assertEqual(response.status_code, 200) + + def _post_request(self): + response = requests.post( + f"{DEFAULT_URL_FOR_TEST}/generate", + json={ + "text": "The capital of France is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + }, + ) + self.assertEqual(response.status_code, 200) + + def _clear_profile_dir(self): + if os.path.isdir(self.output_dir): + shutil.rmtree(self.output_dir) + + def _check_profile_output(self, pattern: str, expect_existence: bool): + self.assertTrue( + os.path.isdir(self.output_dir), "Output directory does not exist." + ) + self.assertEqual( + len(list(Path(self.output_dir).glob(pattern))) > 0, + expect_existence, + f"Does not find {pattern=} ({list(Path(self.output_dir).glob('**/*'))=})", + ) + + +if __name__ == "__main__": + unittest.main()