Move profiler controls to SchedulerProfilerManager (#25614)

python/sglang/srt/managers/scheduler.py

@@ -178,7 +178,6 @@ from sglang.srt.managers.scheduler_output_processor_mixin import (
     SchedulerOutputProcessorMixin,
 )
 from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin
-from sglang.srt.managers.scheduler_profiler_mixin import SchedulerProfilerMixin
 from sglang.srt.managers.scheduler_recv_skipper import SchedulerRecvSkipper
 from sglang.srt.managers.scheduler_runtime_checker_mixin import (
     SchedulerRuntimeCheckerMixin,
@@ -323,7 +322,6 @@ def validate_dflash_request(req: Req) -> Optional[str]:
 class Scheduler(
     SchedulerOutputProcessorMixin,
     SchedulerUpdateWeightsMixin,
-    SchedulerProfilerMixin,
     SchedulerMetricsMixin,
     SchedulerDisaggregationDecodeMixin,
     SchedulerDisaggregationPrefillMixin,
@@ -1325,7 +1323,7 @@ class Scheduler(
                 (SlowDownReqInput, self.slow_down),
                 (
                     ProfileReq,
-                    lambda req: self._profile(self.profiler_manager, req),
+                    lambda req: self.profiler_manager._profile(req),
                 ),
                 (FreezeGCReq, self.handle_freeze_gc),
                 (GetInternalStateReq, self.get_internal_state),
@@ -2725,7 +2723,7 @@ class Scheduler(
         batch.forward_iter = self.forward_ct
 
         # Whether to run the profiler
-        self._profile_batch_predicate(self.profiler_manager, batch)
+        self.profiler_manager._profile_batch_predicate(batch)
         if self.forward_sleep_time is not None:
             logger.info(f"Scheduler.run_batch sleep {self.forward_sleep_time}s")
             time.sleep(self.forward_sleep_time)

python/sglang/srt/managers/scheduler_components/profiler_manager.py

@@ -1,10 +1,44 @@
 from __future__ import annotations
 
+import logging
+import os
+import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, List, Optional
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    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.server_args import get_global_server_args
+from sglang.srt.utils import is_npu
+from sglang.srt.utils.profile_merger import ProfileMerger
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+
+_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__)
+
+
 from sglang.srt.utils.profile_utils import ProfileManager
 
 
@@ -41,3 +75,348 @@ class SchedulerProfilerManager:
 
         # For ROCM
         self.rpd_profiler = None
+
+    def _init_profile(
+        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 = False,
+        profile_prefix: str = "",
+        profile_stages: Optional[List[str]] = None,
+    ) -> 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,
+                profile_stages=profile_stages,
+            )
+
+        if self.profile_in_progress:
+            return ProfileReqOutput(
+                success=False,
+                message="Profiling is already in progress. Call /stop_profile first.",
+            )
+
+        self.profile_by_stage = profile_by_stage
+        self.merge_profiles = merge_profiles
+
+        if output_dir is None:
+            output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")
+        if activities is None:
+            activities = ["CPU", "GPU"]
+
+        self.torch_profiler_output_dir = Path(output_dir).expanduser()
+        self.torch_profiler_with_stack = with_stack
+        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.get_forward_ct() + 1)
+
+        if num_steps:
+            if self.profile_by_stage:
+                self.profiler_prefill_ct = 0
+                self.profiler_decode_ct = 0
+                self.profiler_target_prefill_ct = num_steps
+                self.profiler_target_decode_ct = num_steps
+            elif start_step:
+                self.profiler_target_forward_ct = (
+                    self.profiler_start_forward_ct + num_steps
+                )
+            else:
+                self.profiler_target_forward_ct = self.get_forward_ct() + num_steps
+            # The caller will be notified when reaching profiler_target_forward_ct
+        else:
+            self.profiler_target_forward_ct = None
+
+        return ProfileReqOutput(success=True, message="Succeeded")
+
+    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})",
+        )
+
+        activities = self.profiler_activities
+        with_stack = self.torch_profiler_with_stack
+        record_shapes = self.torch_profiler_record_shapes
+
+        activity_map = {
+            "CPU": torch.profiler.ProfilerActivity.CPU,
+            "GPU": torch.profiler.ProfilerActivity.CUDA,
+        }
+        if hasattr(torch.profiler.ProfilerActivity, "XPU"):
+            activity_map["XPU"] = torch.profiler.ProfilerActivity.XPU
+        torchprof_activities = [
+            activity_map[a] for a in activities if a in activity_map
+        ]
+
+        if "RPD" in activities:  # for ROCM
+            from rpdTracerControl import rpdTracerControl
+
+            rpdTracerControl.skipCreate()
+
+            self.rpd_profile_path = os.path.join(
+                self.torch_profiler_output_dir,
+                "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz",
+            )
+
+            if self.ps.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.dp_tp_cpu_group)
+
+            self.rpd_profiler = rpdTracerControl()
+            self.rpd_profiler.setPythonTrace(True)
+            self.rpd_profiler.start()
+            self.rpd_profiler.rangePush("", "rpd profile range", "")
+            self.profile_in_progress = True
+        elif torchprof_activities:
+            self.torch_profiler = torch.profiler.profile(
+                activities=torchprof_activities,
+                with_stack=with_stack if with_stack is not None else True,
+                record_shapes=record_shapes if record_shapes is not None else False,
+                on_trace_ready=(
+                    None
+                    if not _is_npu
+                    else torch_npu.profiler.tensorboard_trace_handler(
+                        str(self.torch_profiler_output_dir)
+                    )
+                ),
+            )
+            self.torch_profiler.start()
+            self.profile_in_progress = True
+
+        if "MEM" in activities:
+            torch.cuda.memory._record_memory_history(max_entries=100000)
+            self.profile_in_progress = True
+
+        if "CUDA_PROFILER" in activities:
+            if self.ps.gpu_id == get_global_server_args().base_gpu_id:
+                torch.cuda.cudart().cudaProfilerStart()
+            self.profile_in_progress = True
+
+        return ProfileReqOutput(success=True, message="Succeeded")
+
+    def _merge_profile_traces(self) -> str:
+        if not self.merge_profiles:
+            return ""
+
+        if self.ps.tp_rank != 0:
+            return ""
+        if self.ps.dp_size > 1 and self.ps.dp_rank != 0:
+            return ""
+        if self.ps.pp_size > 1 and self.ps.pp_rank != 0:
+            return ""
+        if self.ps.moe_ep_size > 1 and self.ps.moe_ep_rank != 0:
+            return ""
+
+        try:
+            logger.info("Starting profile merge...")
+            merger = ProfileMerger(self.torch_profiler_output_dir, self.profile_id)
+            merged_path = merger.merge_chrome_traces()
+
+            summary = merger.get_merge_summary()
+            merge_message = (
+                f" Merged trace: {merged_path} "
+                f"(Events: {summary.get('total_events', '?')}, "
+                f"Files: {summary.get('total_files', '?')})"
+            )
+
+            logger.info(f"Profile merge completed: {merged_path}")
+        except Exception as e:
+            logger.error(f"Failed to merge profiles: {e}", exc_info=True)
+            return f" Merge failed: {e!s}"
+        else:
+            return merge_message
+
+    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,
+                message="Profiling is not in progress. Call /start_profile first.",
+            )
+
+        self.torch_profiler_output_dir.mkdir(parents=True, exist_ok=True)
+
+        if self.profile_prefix:
+            stage_prefix = self.profile_prefix + "-"
+        else:
+            stage_prefix = ""
+
+        stage_suffix = f"-{stage.name}" if stage else ""
+        logger.info("Stop profiling" + stage_suffix + "...")
+        if self.torch_profiler is not None:
+            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.ps.tp_rank}"]
+
+                # Only add other ranks if parallelism is enabled (size > 1)
+                if self.ps.dp_size > 1:
+                    filename_parts.append(f"DP-{self.ps.dp_rank}")
+                if self.ps.pp_size > 1:
+                    filename_parts.append(f"PP-{self.ps.pp_rank}")
+                if self.ps.moe_ep_size > 1:
+                    filename_parts.append(f"EP-{self.ps.moe_ep_rank}")
+
+                filename = (
+                    stage_prefix
+                    + "-".join(filename_parts)
+                    + stage_suffix
+                    + ".trace.json.gz"
+                )
+
+                self.torch_profiler.export_chrome_trace(
+                    os.path.join(self.torch_profiler_output_dir, filename)
+                )
+            torch.distributed.barrier(self.dp_tp_cpu_group)
+
+        if self.rpd_profiler is not None:
+            self.rpd_profiler.rangePop()
+            self.rpd_profiler.stop()
+            self.rpd_profiler.flush()
+
+            torch.distributed.barrier(self.dp_tp_cpu_group)
+            if self.ps.tp_rank == 0:
+                from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace
+
+                rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
+            self.rpd_profiler = None
+            self.rpd_profile_path = None
+
+        if self.profiler_activities is not None and "MEM" in self.profiler_activities:
+            memory_profile_path = os.path.join(
+                self.torch_profiler_output_dir,
+                str(time.time())
+                + f"-TP-{self.ps.tp_rank}-memory"
+                + stage_suffix
+                + ".pickle",
+            )
+            torch.cuda.memory._dump_snapshot(memory_profile_path)
+            torch.cuda.memory._record_memory_history(enabled=None)
+
+        if "CUDA_PROFILER" in self.profiler_activities:
+            if self.ps.gpu_id == get_global_server_args().base_gpu_id:
+                torch.cuda.cudart().cudaProfilerStop()
+
+        merge_message = self._merge_profile_traces()
+
+        logger.info(
+            "Profiling done. Traces are saved to: %s%s",
+            self.torch_profiler_output_dir,
+            merge_message,
+        )
+        self.torch_profiler = None
+        self.profile_in_progress = False
+        self.profiler_start_forward_ct = None
+
+        return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}")
+
+    def _profile_batch_predicate(self, batch: ScheduleBatch):
+        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:
+                    self._start_profile(batch.forward_mode)
+                self.profiler_prefill_ct += 1
+                if self.profiler_prefill_ct > self.profiler_target_prefill_ct:
+                    if self.profile_in_progress:
+                        self._stop_profile(stage=ForwardMode.EXTEND)
+            elif batch.forward_mode.is_decode():
+                if self.profiler_decode_ct == 0:
+                    if self.profile_in_progress:
+                        # force trace flush
+                        self._stop_profile(stage=ForwardMode.EXTEND)
+                    self._start_profile(batch.forward_mode)
+                self.profiler_decode_ct += 1
+                if self.profiler_decode_ct > self.profiler_target_decode_ct:
+                    if self.profile_in_progress:
+                        self._stop_profile(stage=ForwardMode.DECODE)
+            elif batch.forward_mode.is_idle():
+                pass
+            else:
+                raise RuntimeError(f"unsupported profile stage: {batch.forward_mode}")
+        else:
+            # Check profiler
+            if (
+                self.profiler_target_forward_ct
+                and self.profiler_target_forward_ct <= self.get_forward_ct()
+            ):
+                self._stop_profile()
+            if (
+                self.profiler_start_forward_ct
+                and self.profiler_start_forward_ct == self.get_forward_ct()
+            ):
+                self._start_profile()
+
+    def _profile(self, recv_req: ProfileReq):
+        if recv_req.type == ProfileReqType.START_PROFILE:
+            if recv_req.profile_by_stage or recv_req.start_step:
+                return self._init_profile(
+                    recv_req.output_dir,
+                    recv_req.start_step,
+                    recv_req.num_steps,
+                    recv_req.activities,
+                    recv_req.with_stack,
+                    recv_req.record_shapes,
+                    recv_req.profile_by_stage,
+                    recv_req.profile_id,
+                    recv_req.merge_profiles,
+                    recv_req.profile_prefix,
+                    recv_req.profile_stages,
+                )
+            else:
+                self._init_profile(
+                    recv_req.output_dir,
+                    recv_req.start_step,
+                    recv_req.num_steps,
+                    recv_req.activities,
+                    recv_req.with_stack,
+                    recv_req.record_shapes,
+                    recv_req.profile_by_stage,
+                    recv_req.profile_id,
+                    recv_req.merge_profiles,
+                    recv_req.profile_prefix,
+                )
+                return self._start_profile()
+        else:
+            return self._stop_profile()

python/sglang/srt/managers/scheduler_profiler_mixin.py

@@ -1,398 +0,0 @@
-from __future__ import annotations
-
-import logging
-import os
-import time
-from pathlib import Path
-from typing import TYPE_CHECKING, 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.server_args import get_global_server_args
-from sglang.srt.utils import is_npu
-from sglang.srt.utils.profile_merger import ProfileMerger
-
-if TYPE_CHECKING:
-    from sglang.srt.managers.schedule_batch import ScheduleBatch
-    from sglang.srt.managers.scheduler_components.profiler_manager import (
-        SchedulerProfilerManager,
-    )
-
-_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 SchedulerProfilerMixin:
-    @staticmethod
-    def _init_profile(
-        self: "SchedulerProfilerManager",
-        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 = False,
-        profile_prefix: str = "",
-        profile_stages: Optional[List[str]] = None,
-    ) -> 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,
-                profile_stages=profile_stages,
-            )
-
-        if self.profile_in_progress:
-            return ProfileReqOutput(
-                success=False,
-                message="Profiling is already in progress. Call /stop_profile first.",
-            )
-
-        self.profile_by_stage = profile_by_stage
-        self.merge_profiles = merge_profiles
-
-        if output_dir is None:
-            output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")
-        if activities is None:
-            activities = ["CPU", "GPU"]
-
-        self.torch_profiler_output_dir = Path(output_dir).expanduser()
-        self.torch_profiler_with_stack = with_stack
-        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.get_forward_ct() + 1)
-
-        if num_steps:
-            if self.profile_by_stage:
-                self.profiler_prefill_ct = 0
-                self.profiler_decode_ct = 0
-                self.profiler_target_prefill_ct = num_steps
-                self.profiler_target_decode_ct = num_steps
-            elif start_step:
-                self.profiler_target_forward_ct = (
-                    self.profiler_start_forward_ct + num_steps
-                )
-            else:
-                self.profiler_target_forward_ct = self.get_forward_ct() + num_steps
-            # The caller will be notified when reaching profiler_target_forward_ct
-        else:
-            self.profiler_target_forward_ct = None
-
-        return ProfileReqOutput(success=True, message="Succeeded")
-
-    @staticmethod
-    def _start_profile(
-        self: "SchedulerProfilerManager", 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})",
-        )
-
-        activities = self.profiler_activities
-        with_stack = self.torch_profiler_with_stack
-        record_shapes = self.torch_profiler_record_shapes
-
-        activity_map = {
-            "CPU": torch.profiler.ProfilerActivity.CPU,
-            "GPU": torch.profiler.ProfilerActivity.CUDA,
-        }
-        if hasattr(torch.profiler.ProfilerActivity, "XPU"):
-            activity_map["XPU"] = torch.profiler.ProfilerActivity.XPU
-        torchprof_activities = [
-            activity_map[a] for a in activities if a in activity_map
-        ]
-
-        if "RPD" in activities:  # for ROCM
-            from rpdTracerControl import rpdTracerControl
-
-            rpdTracerControl.skipCreate()
-
-            self.rpd_profile_path = os.path.join(
-                self.torch_profiler_output_dir,
-                "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz",
-            )
-
-            if self.ps.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.dp_tp_cpu_group)
-
-            self.rpd_profiler = rpdTracerControl()
-            self.rpd_profiler.setPythonTrace(True)
-            self.rpd_profiler.start()
-            self.rpd_profiler.rangePush("", "rpd profile range", "")
-            self.profile_in_progress = True
-        elif torchprof_activities:
-            self.torch_profiler = torch.profiler.profile(
-                activities=torchprof_activities,
-                with_stack=with_stack if with_stack is not None else True,
-                record_shapes=record_shapes if record_shapes is not None else False,
-                on_trace_ready=(
-                    None
-                    if not _is_npu
-                    else torch_npu.profiler.tensorboard_trace_handler(
-                        str(self.torch_profiler_output_dir)
-                    )
-                ),
-            )
-            self.torch_profiler.start()
-            self.profile_in_progress = True
-
-        if "MEM" in activities:
-            torch.cuda.memory._record_memory_history(max_entries=100000)
-            self.profile_in_progress = True
-
-        if "CUDA_PROFILER" in activities:
-            if self.ps.gpu_id == get_global_server_args().base_gpu_id:
-                torch.cuda.cudart().cudaProfilerStart()
-            self.profile_in_progress = True
-
-        return ProfileReqOutput(success=True, message="Succeeded")
-
-    @staticmethod
-    def _merge_profile_traces(self: "SchedulerProfilerManager") -> str:
-        if not self.merge_profiles:
-            return ""
-
-        if self.ps.tp_rank != 0:
-            return ""
-        if self.ps.dp_size > 1 and self.ps.dp_rank != 0:
-            return ""
-        if self.ps.pp_size > 1 and self.ps.pp_rank != 0:
-            return ""
-        if self.ps.moe_ep_size > 1 and self.ps.moe_ep_rank != 0:
-            return ""
-
-        try:
-            logger.info("Starting profile merge...")
-            merger = ProfileMerger(self.torch_profiler_output_dir, self.profile_id)
-            merged_path = merger.merge_chrome_traces()
-
-            summary = merger.get_merge_summary()
-            merge_message = (
-                f" Merged trace: {merged_path} "
-                f"(Events: {summary.get('total_events', '?')}, "
-                f"Files: {summary.get('total_files', '?')})"
-            )
-
-            logger.info(f"Profile merge completed: {merged_path}")
-        except Exception as e:
-            logger.error(f"Failed to merge profiles: {e}", exc_info=True)
-            return f" Merge failed: {e!s}"
-        else:
-            return merge_message
-
-    @staticmethod
-    def _stop_profile(
-        self: "SchedulerProfilerManager", 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,
-                message="Profiling is not in progress. Call /start_profile first.",
-            )
-
-        self.torch_profiler_output_dir.mkdir(parents=True, exist_ok=True)
-
-        if self.profile_prefix:
-            stage_prefix = self.profile_prefix + "-"
-        else:
-            stage_prefix = ""
-
-        stage_suffix = f"-{stage.name}" if stage else ""
-        logger.info("Stop profiling" + stage_suffix + "...")
-        if self.torch_profiler is not None:
-            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.ps.tp_rank}"]
-
-                # Only add other ranks if parallelism is enabled (size > 1)
-                if self.ps.dp_size > 1:
-                    filename_parts.append(f"DP-{self.ps.dp_rank}")
-                if self.ps.pp_size > 1:
-                    filename_parts.append(f"PP-{self.ps.pp_rank}")
-                if self.ps.moe_ep_size > 1:
-                    filename_parts.append(f"EP-{self.ps.moe_ep_rank}")
-
-                filename = (
-                    stage_prefix
-                    + "-".join(filename_parts)
-                    + stage_suffix
-                    + ".trace.json.gz"
-                )
-
-                self.torch_profiler.export_chrome_trace(
-                    os.path.join(self.torch_profiler_output_dir, filename)
-                )
-            torch.distributed.barrier(self.dp_tp_cpu_group)
-
-        if self.rpd_profiler is not None:
-            self.rpd_profiler.rangePop()
-            self.rpd_profiler.stop()
-            self.rpd_profiler.flush()
-
-            torch.distributed.barrier(self.dp_tp_cpu_group)
-            if self.ps.tp_rank == 0:
-                from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace
-
-                rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
-            self.rpd_profiler = None
-            self.rpd_profile_path = None
-
-        if self.profiler_activities is not None and "MEM" in self.profiler_activities:
-            memory_profile_path = os.path.join(
-                self.torch_profiler_output_dir,
-                str(time.time())
-                + f"-TP-{self.ps.tp_rank}-memory"
-                + stage_suffix
-                + ".pickle",
-            )
-            torch.cuda.memory._dump_snapshot(memory_profile_path)
-            torch.cuda.memory._record_memory_history(enabled=None)
-
-        if "CUDA_PROFILER" in self.profiler_activities:
-            if self.ps.gpu_id == get_global_server_args().base_gpu_id:
-                torch.cuda.cudart().cudaProfilerStop()
-
-        merge_message = SchedulerProfilerMixin._merge_profile_traces(self)
-
-        logger.info(
-            "Profiling done. Traces are saved to: %s%s",
-            self.torch_profiler_output_dir,
-            merge_message,
-        )
-        self.torch_profiler = None
-        self.profile_in_progress = False
-        self.profiler_start_forward_ct = None
-
-        return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}")
-
-    @staticmethod
-    def _profile_batch_predicate(
-        self: "SchedulerProfilerManager", batch: ScheduleBatch
-    ):
-        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:
-                    SchedulerProfilerMixin._start_profile(self, batch.forward_mode)
-                self.profiler_prefill_ct += 1
-                if self.profiler_prefill_ct > self.profiler_target_prefill_ct:
-                    if self.profile_in_progress:
-                        SchedulerProfilerMixin._stop_profile(
-                            self, stage=ForwardMode.EXTEND
-                        )
-            elif batch.forward_mode.is_decode():
-                if self.profiler_decode_ct == 0:
-                    if self.profile_in_progress:
-                        # force trace flush
-                        SchedulerProfilerMixin._stop_profile(
-                            self, stage=ForwardMode.EXTEND
-                        )
-                    SchedulerProfilerMixin._start_profile(self, batch.forward_mode)
-                self.profiler_decode_ct += 1
-                if self.profiler_decode_ct > self.profiler_target_decode_ct:
-                    if self.profile_in_progress:
-                        SchedulerProfilerMixin._stop_profile(
-                            self, stage=ForwardMode.DECODE
-                        )
-            elif batch.forward_mode.is_idle():
-                pass
-            else:
-                raise RuntimeError(f"unsupported profile stage: {batch.forward_mode}")
-        else:
-            # Check profiler
-            if (
-                self.profiler_target_forward_ct
-                and self.profiler_target_forward_ct <= self.get_forward_ct()
-            ):
-                SchedulerProfilerMixin._stop_profile(self)
-            if (
-                self.profiler_start_forward_ct
-                and self.profiler_start_forward_ct == self.get_forward_ct()
-            ):
-                SchedulerProfilerMixin._start_profile(self)
-
-    @staticmethod
-    def _profile(self: "SchedulerProfilerManager", recv_req: ProfileReq):
-        if recv_req.type == ProfileReqType.START_PROFILE:
-            if recv_req.profile_by_stage or recv_req.start_step:
-                return SchedulerProfilerMixin._init_profile(
-                    self,
-                    recv_req.output_dir,
-                    recv_req.start_step,
-                    recv_req.num_steps,
-                    recv_req.activities,
-                    recv_req.with_stack,
-                    recv_req.record_shapes,
-                    recv_req.profile_by_stage,
-                    recv_req.profile_id,
-                    recv_req.merge_profiles,
-                    recv_req.profile_prefix,
-                    recv_req.profile_stages,
-                )
-            else:
-                SchedulerProfilerMixin._init_profile(
-                    self,
-                    recv_req.output_dir,
-                    recv_req.start_step,
-                    recv_req.num_steps,
-                    recv_req.activities,
-                    recv_req.with_stack,
-                    recv_req.record_shapes,
-                    recv_req.profile_by_stage,
-                    recv_req.profile_id,
-                    recv_req.merge_profiles,
-                    recv_req.profile_prefix,
-                )
-                return SchedulerProfilerMixin._start_profile(self)
-        else:
-            return SchedulerProfilerMixin._stop_profile(self)

test/registered/unit/utils/test_profile_merger.py

@@ -229,9 +229,11 @@ class TestProfileMergerIntegration(unittest.TestCase):
         self.assertIn("merge_profiles", sig.parameters)
 
         # Test SchedulerProfilerMixin
-        from sglang.srt.managers.scheduler_profiler_mixin import SchedulerProfilerMixin
+        from sglang.srt.managers.scheduler_components.profiler_manager import (
+            SchedulerProfilerManager,
+        )
 
-        sig = inspect.signature(SchedulerProfilerMixin._init_profile)
+        sig = inspect.signature(SchedulerProfilerManager._init_profile)
         self.assertIn("merge_profiles", sig.parameters)
 
         # Test CLI profiler