diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 1ebeeec4f..f1c65a3b1 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py @@ -18,6 +18,7 @@ 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.platforms import current_platform from sglang.srt.runtime_context import get_server_args from sglang.srt.utils import is_mps, is_npu from sglang.srt.utils.profile_merger import ProfileMerger @@ -170,6 +171,15 @@ class SchedulerProfilerManager: "CPU": torch.profiler.ProfilerActivity.CPU, "GPU": torch.profiler.ProfilerActivity.CUDA, } + + if current_platform.is_out_of_tree(): + if hasattr( + torch.profiler.ProfilerActivity, + current_platform.get_torch_profiler_activity_str(), + ): + activity_map[current_platform.get_torch_profiler_activity_str()] = ( + current_platform.get_torch_profiler_activity() + ) if hasattr(torch.profiler.ProfilerActivity, "XPU"): activity_map["XPU"] = torch.profiler.ProfilerActivity.XPU torchprof_activities = [ diff --git a/python/sglang/srt/platforms/device_mixin.py b/python/sglang/srt/platforms/device_mixin.py index 5781b0191..fd51a24ae 100644 --- a/python/sglang/srt/platforms/device_mixin.py +++ b/python/sglang/srt/platforms/device_mixin.py @@ -251,6 +251,14 @@ class DeviceMixin: return CpuArchEnum.ARM return CpuArchEnum.UNSPECIFIED + def get_torch_profiler_activity_str(self) -> str: + """[Planned] Return the torch profiler activity string.""" + raise NotImplementedError + + def get_torch_profiler_activity(self) -> torch.profiler.ProfilerActivity: + """[Planned] Return the torch profiler activity.""" + raise NotImplementedError + # ------------------------------------------------------------------ # Dunder helpers # ------------------------------------------------------------------ diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index e4387bd1c..693a7418d 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/python/sglang/srt/utils/profile_utils.py @@ -12,6 +12,7 @@ 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 ForwardBatch, ForwardMode +from sglang.srt.platforms import current_platform 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 @@ -221,6 +222,16 @@ class _ProfilerBase(ABC): @staticmethod def create(activities, with_stack, record_shapes, **kwargs): inners = [] + if current_platform.is_out_of_tree(): + if current_platform.get_torch_profiler_activity_str() in activities: + inners.append( + _ProfilerTorch( + **kwargs, + activities=activities, + with_stack=with_stack, + record_shapes=record_shapes, + ) + ) if ("CPU" in activities) or ("GPU" in activities): inners.append( _ProfilerTorch( @@ -291,6 +302,12 @@ class _ProfilerTorch(_ProfilerConcreteBase): "CPU": torch.profiler.ProfilerActivity.CPU, "GPU": torch.profiler.ProfilerActivity.CUDA, } + + if current_platform.is_out_of_tree(): + activity_map[current_platform.get_torch_profiler_activity_str()] = ( + current_platform.get_torch_profiler_activity() + ) + torchprof_activities = [ activity_map[a] for a in self.activities if a in activity_map ]