[plugin] oot torch profiler activity support (#31580)

Signed-off-by: Devashish Lal <devcode@fb.com>
Co-authored-by: Devashish Lal <devcode@fb.com>

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 = [

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
     # ------------------------------------------------------------------

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
         ]