Implement profiler v2 and fix stage mixture bug (#14148)

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:

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)

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"),
     ],
 }
 

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()