chore: import upstream snapshot with attribution

python/sglang/srt/utils/profile_utils.py

@@ -0,0 +1,414 @@
+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.distributed.parallel_state_wrapper import ParallelState
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import ProfileReqOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils import is_npu
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
+
+_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],
+    ]
+    apply_torch_npu_patches(torch_npu, patches)
+
+logger = logging.getLogger(__name__)
+
+
+def export_cuda_graph_capture_trace(prof_context, *, runner_name: str, tp_rank: int):
+    """Persist a CUDA-graph capture profiler trace (chrome trace) to disk.
+
+    Opt-in via ``SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE`` (no-op otherwise). The
+    capture profiler must have run with ``record_shapes=True`` so the trace can
+    be inspected offline as a per-kernel shape/identity record. The file lands in
+    ``<SGLANG_TORCH_PROFILER_DIR>/graph_capture_profile/`` and is namespaced by
+    runner class and TP rank so concurrent capture passes (e.g. EAGLE3
+    target/draft/draft-extend) and ranks don't overwrite each other.
+    """
+    if not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get():
+        return
+    output_dir = os.path.join(
+        envs.SGLANG_TORCH_PROFILER_DIR.get(), "graph_capture_profile"
+    )
+    os.makedirs(output_dir, exist_ok=True)
+    path = os.path.join(
+        output_dir, f"cuda_graph_capture-{runner_name}-TP-{tp_rank}.json.gz"
+    )
+    prof_context.export_chrome_trace(path)
+    logger.info(f"CUDA graph capture trace saved to: {path}")
+
+
+class ProfileManager:
+    def __init__(self, ps: ParallelState, cpu_group):
+        self.stage_based_trigger = _StageBasedTrigger(
+            on_start=self._do_start,
+            on_stop=self._do_stop,
+        )
+        self.ps = ps
+        self.cpu_group = cpu_group
+        self.first_rank_in_node = ps.gpu_id == get_server_args().base_gpu_id
+        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,
+        profile_stages: Optional[List[str]] = None,
+    ):
+        # 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=profile_stages or ["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,
+            ps=self.ps,
+            cpu_group=self.cpu_group,
+            first_rank_in_node=self.first_rank_in_node,
+            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,
+        ps: ParallelState,
+        cpu_group,
+        first_rank_in_node: bool,
+    ):
+        self.output_dir = output_dir
+        self.output_prefix = output_prefix
+        self.output_suffix = output_suffix
+        self.profile_id = profile_id
+        self.ps = ps
+        self.cpu_group = cpu_group
+        self.first_rank_in_node = first_rank_in_node
+
+
+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.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 = (
+                (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.ps.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):
+        if self.first_rank_in_node:
+            logger.info(f"Call cudaProfilerStart")
+            torch.cuda.cudart().cudaProfilerStart()
+
+    def stop(self):
+        if self.first_rank_in_node:
+            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.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.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.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)