[MLX] Add Metal profiling hooks to server profiler (#28122)

python/sglang/srt/hardware_backend/mlx/profiler.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import gzip
+import json
+import logging
+import os
+import shutil
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Optional
+
+import torch
+
+from sglang.srt.managers.io_struct import ProfileReqOutput
+from sglang.srt.utils.tensor_bridge import use_mlx
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MetalCaptureProfiler:
+    label: str
+    trace_path: Path
+    stop_capture: Callable[[], None]
+    standalone: bool
+
+    @classmethod
+    def start_mlx(cls, trace_path: Path):
+        trace_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            import mlx.core as mx
+
+            mx.metal.start_capture(str(trace_path))
+        except RuntimeError as e:
+            return None, _capture_error("MLX", e)
+
+        return cls._started(
+            label="MLX",
+            trace_path=trace_path,
+            stop_capture=mx.metal.stop_capture,
+            standalone=True,
+        )
+
+    @classmethod
+    def start_mps(cls, trace_path: Path):
+        trace_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            if not hasattr(torch, "mps") or not hasattr(torch.mps, "profiler"):
+                raise RuntimeError("torch.mps.profiler is not available")
+            context = torch.mps.profiler.metal_capture(str(trace_path))
+            context.__enter__()
+        except RuntimeError as e:
+            return None, _capture_error("MPS", e)
+
+        return cls._started(
+            label="MPS",
+            trace_path=trace_path,
+            stop_capture=lambda: context.__exit__(None, None, None),
+            standalone=False,
+        )
+
+    @classmethod
+    def _started(
+        cls,
+        *,
+        label: str,
+        trace_path: Path,
+        stop_capture: Callable[[], None],
+        standalone: bool,
+    ):
+        profiler = cls(
+            label=label,
+            trace_path=trace_path,
+            stop_capture=stop_capture,
+            standalone=standalone,
+        )
+        logger.info("%s Metal capture started, saving to %s", label, trace_path)
+        return profiler, ProfileReqOutput(success=True, message="Succeeded")
+
+    def stop(self) -> str:
+        self.stop_capture()
+
+        logger.info(
+            "%s Metal capture stopped. Trace saved to: %s",
+            self.label,
+            self.trace_path,
+        )
+        return f" Metal trace: {self.trace_path}"
+
+
+def _capture_error(label: str, error: RuntimeError) -> ProfileReqOutput:
+    return ProfileReqOutput(
+        success=False,
+        message=(
+            f"Failed to start {label} Metal capture: {error}. "
+            "Set MTL_CAPTURE_ENABLED=1 in the server's environment "
+            "before launching to enable GPU trace capture."
+        ),
+    )
+
+
+class MetalTorchProfiler:
+    def __init__(
+        self,
+        *,
+        start_metal_capture: Callable[[Path], tuple[Any, ProfileReqOutput]],
+        torch_profiler: Optional[Any] = None,
+    ):
+        self.start_metal_capture = start_metal_capture
+        self.torch_profiler = torch_profiler
+        self.metal_profiler = None
+
+    def start(self):
+        trace_path = _new_temp_gputrace_path()
+        self.metal_profiler, result = self.start_metal_capture(trace_path)
+        if not result.success:
+            raise RuntimeError(result.message)
+        if self.torch_profiler is not None:
+            try:
+                self.torch_profiler.start()
+            except Exception:
+                self.metal_profiler.stop()
+                raise
+
+    def stop(self):
+        try:
+            if self.torch_profiler is not None:
+                self.torch_profiler.stop()
+        finally:
+            if self.metal_profiler is not None:
+                self.metal_profiler.stop()
+
+    def export_chrome_trace(self, path: str):
+        if self.torch_profiler is not None:
+            self.torch_profiler.export_chrome_trace(path)
+        else:
+            _write_empty_chrome_trace(path)
+
+        if self.metal_profiler is None:
+            return
+
+        final_path = _unique_gputrace_path_for_chrome_trace(path)
+        final_path.parent.mkdir(parents=True, exist_ok=True)
+        if self.metal_profiler.trace_path.exists():
+            shutil.move(str(self.metal_profiler.trace_path), str(final_path))
+            logger.info("Metal trace saved to: %s", final_path)
+
+
+def apply_metal_profiler_patches() -> None:
+    if getattr(torch.profiler.profile, "_sglang_metal_patched", False):
+        return
+
+    original_profile = torch.profiler.profile
+
+    def profile(*args, **kwargs):
+        activities = _get_activities(args, kwargs)
+        if not _has_cuda_activity(activities):
+            return original_profile(*args, **kwargs)
+
+        if use_mlx():
+            return MetalTorchProfiler(
+                start_metal_capture=MetalCaptureProfiler.start_mlx
+            )
+
+        torch_activities = [
+            activity for activity in activities if not _is_cuda_activity(activity)
+        ]
+        torch_profiler = None
+        if torch_activities:
+            patched_args, patched_kwargs = _replace_activities(
+                args, kwargs, torch_activities
+            )
+            torch_profiler = original_profile(*patched_args, **patched_kwargs)
+
+        return MetalTorchProfiler(
+            start_metal_capture=MetalCaptureProfiler.start_mps,
+            torch_profiler=torch_profiler,
+        )
+
+    profile._sglang_metal_patched = True
+    profile._sglang_original_profile = original_profile
+    torch.profiler.profile = profile
+
+
+def _get_activities(args, kwargs):
+    if "activities" in kwargs:
+        return kwargs["activities"]
+    if args:
+        return args[0]
+    return None
+
+
+def _replace_activities(args, kwargs, activities):
+    kwargs = dict(kwargs)
+    if "activities" in kwargs:
+        kwargs["activities"] = activities
+        return args, kwargs
+
+    if args:
+        args = list(args)
+        args[0] = activities
+        return tuple(args), kwargs
+
+    kwargs["activities"] = activities
+    return args, kwargs
+
+
+def _has_cuda_activity(activities) -> bool:
+    if activities is None:
+        return False
+    return any(_is_cuda_activity(activity) for activity in activities)
+
+
+def _is_cuda_activity(activity) -> bool:
+    return activity == torch.profiler.ProfilerActivity.CUDA
+
+
+def _new_temp_gputrace_path() -> Path:
+    output_dir = Path(os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")).expanduser()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for i in range(100):
+        candidate = (
+            output_dir / f"sglang-metal-{os.getpid()}-{time.time_ns()}-{i}.gputrace"
+        )
+        if not candidate.exists():
+            return candidate
+    raise RuntimeError(f"Cannot find an unused Metal trace path in {output_dir}")
+
+
+def _unique_gputrace_path_for_chrome_trace(path: str) -> Path:
+    chrome_path = Path(path).expanduser()
+    name = chrome_path.name
+    if name.endswith(".trace.json.gz"):
+        name = name[: -len(".trace.json.gz")] + ".gputrace"
+    else:
+        name = chrome_path.stem + ".gputrace"
+
+    base = chrome_path.with_name(name)
+    if not base.exists():
+        return base
+
+    stem = base.name[: -len(".gputrace")]
+    for i in range(100):
+        candidate = base.with_name(f"{stem}-{time.time_ns()}-{i}.gputrace")
+        if not candidate.exists():
+            return candidate
+    raise RuntimeError(f"Cannot find an unused Metal trace path for {base}")
+
+
+def _write_empty_chrome_trace(path: str):
+    trace = {"traceEvents": []}
+    Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
+    if str(path).endswith(".gz"):
+        with gzip.open(path, "wt") as f:
+            json.dump(trace, f)
+    else:
+        with open(path, "w") as f:
+            json.dump(trace, f)

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

@@ -19,14 +19,16 @@ 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 import is_mps, is_npu
 from sglang.srt.utils.profile_merger import ProfileMerger
+from sglang.srt.utils.profile_utils import ProfileManager
 from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
 
 _is_npu = is_npu()
+_is_mps = is_mps()
 if _is_npu:
     import torch_npu
 
@@ -36,13 +38,14 @@ if _is_npu:
         ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU],
     ]
     apply_torch_npu_patches(torch_npu, patches)
+elif _is_mps:
+    from sglang.srt.hardware_backend.mlx.profiler import apply_metal_profiler_patches
+
+    apply_metal_profiler_patches()
 
 logger = logging.getLogger(__name__)
 
 
-from sglang.srt.utils.profile_utils import ProfileManager
-
-
 @dataclass(kw_only=True)
 class SchedulerProfilerManager:
     ps: Any
@@ -230,7 +233,11 @@ class SchedulerProfilerManager:
                     )
                 ),
             )
-            self.torch_profiler.start()
+            try:
+                self.torch_profiler.start()
+            except RuntimeError as e:
+                self.torch_profiler = None
+                return ProfileReqOutput(success=False, message=str(e))
             self.profile_in_progress = True
 
         if "MEM" in activities:

test/registered/unit/hardware_backend/mlx/test_metal_profiler.py

@@ -0,0 +1,251 @@
+"""Unit tests for MLX Metal profiling patch (hardware_backend/mlx/profiler.py).
+
+Covers:
+  - apply_metal_profiler_patches() replaces torch.profiler.profile
+  - MLX path: MetalTorchProfiler.start/stop produces a .gputrace file
+  - MPS path: MetalTorchProfiler wraps torch.mps.profiler.metal_capture
+  - RuntimeError from start_capture is caught and returned as success=False
+  - SchedulerProfilerManager._start_profile returns success=False gracefully
+    when Metal capture fails (no MTL_CAPTURE_ENABLED)
+
+Skips on non-Apple-Silicon platforms and when ``mlx`` is missing.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import platform
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
+_HAS_MLX = importlib.util.find_spec("mlx") is not None
+_SKIP_REASON = "requires Apple Silicon and mlx"
+
+
+@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
+class TestApplyMetalProfilerPatches(unittest.TestCase):
+    """apply_metal_profiler_patches() replaces torch.profiler.profile."""
+
+    def setUp(self):
+        import torch
+
+        self._original_profile = getattr(
+            torch.profiler.profile, "_sglang_original_profile", None
+        )
+
+    def tearDown(self):
+        import torch
+
+        if self._original_profile is not None:
+            torch.profiler.profile = self._original_profile
+
+    def test_patch_replaces_profile(self):
+        import torch
+
+        from sglang.srt.hardware_backend.mlx.profiler import (
+            MetalTorchProfiler,
+            apply_metal_profiler_patches,
+        )
+
+        apply_metal_profiler_patches()
+        self.assertTrue(getattr(torch.profiler.profile, "_sglang_metal_patched", False))
+        p = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA])
+        self.assertIsInstance(p, MetalTorchProfiler)
+
+    def test_patch_is_idempotent(self):
+        import torch
+
+        from sglang.srt.hardware_backend.mlx.profiler import (
+            apply_metal_profiler_patches,
+        )
+
+        apply_metal_profiler_patches()
+        first = torch.profiler.profile
+        apply_metal_profiler_patches()
+        self.assertIs(torch.profiler.profile, first)
+
+    def test_no_cuda_activity_uses_original(self):
+        import torch
+
+        from sglang.srt.hardware_backend.mlx.profiler import (
+            MetalTorchProfiler,
+            apply_metal_profiler_patches,
+        )
+
+        apply_metal_profiler_patches()
+        p = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
+        self.assertNotIsInstance(p, MetalTorchProfiler)
+
+
+@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
+class TestMetalCaptureProfilerMLX(unittest.TestCase):
+    """MLX path: start_mlx produces a .gputrace and stop_capture is called."""
+
+    def test_start_mlx_success(self):
+        import mlx.core as mx
+
+        from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
+
+        with tempfile.TemporaryDirectory() as tmp:
+            trace_path = Path(tmp) / "test.gputrace"
+            with patch.object(mx.metal, "start_capture"), patch.object(
+                mx.metal, "stop_capture"
+            ):
+                profiler, result = MetalCaptureProfiler.start_mlx(trace_path)
+
+            self.assertTrue(result.success)
+            self.assertIsNotNone(profiler)
+            self.assertEqual(profiler.label, "MLX")
+            self.assertTrue(profiler.standalone)
+
+    def test_start_mlx_runtime_error_returns_failure(self):
+        import mlx.core as mx
+
+        from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
+
+        with tempfile.TemporaryDirectory() as tmp:
+            trace_path = Path(tmp) / "test.gputrace"
+            with patch.object(
+                mx.metal,
+                "start_capture",
+                side_effect=RuntimeError("Capture layer is not inserted"),
+            ):
+                profiler, result = MetalCaptureProfiler.start_mlx(trace_path)
+
+        self.assertIsNone(profiler)
+        self.assertFalse(result.success)
+        self.assertIn("MTL_CAPTURE_ENABLED", result.message)
+
+    def test_stop_calls_stop_capture(self):
+        import mlx.core as mx
+
+        from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
+
+        with tempfile.TemporaryDirectory() as tmp:
+            trace_path = Path(tmp) / "test.gputrace"
+            with patch.object(mx.metal, "start_capture"), patch.object(
+                mx.metal, "stop_capture"
+            ) as mock_stop:
+                profiler, _ = MetalCaptureProfiler.start_mlx(trace_path)
+                profiler.stop()
+                mock_stop.assert_called_once()
+
+
+@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
+class TestMetalCaptureProfilerMPS(unittest.TestCase):
+    """MPS path: start_mps wraps torch.mps.profiler.metal_capture."""
+
+    def test_start_mps_success(self):
+        import torch
+
+        from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
+
+        mock_ctx = MagicMock()
+        mock_ctx.__enter__ = MagicMock(return_value=mock_ctx)
+        mock_ctx.__exit__ = MagicMock(return_value=False)
+
+        with tempfile.TemporaryDirectory() as tmp:
+            trace_path = Path(tmp) / "test.gputrace"
+            with patch.object(
+                torch.mps.profiler, "metal_capture", return_value=mock_ctx
+            ):
+                profiler, result = MetalCaptureProfiler.start_mps(trace_path)
+
+        self.assertTrue(result.success)
+        self.assertIsNotNone(profiler)
+        self.assertEqual(profiler.label, "MPS")
+        self.assertFalse(profiler.standalone)
+
+    def test_start_mps_runtime_error_returns_failure(self):
+        import torch
+
+        from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
+
+        with tempfile.TemporaryDirectory() as tmp:
+            trace_path = Path(tmp) / "test.gputrace"
+            with patch.object(
+                torch.mps.profiler,
+                "metal_capture",
+                side_effect=RuntimeError("MPS profiler unavailable"),
+            ):
+                profiler, result = MetalCaptureProfiler.start_mps(trace_path)
+
+        self.assertIsNone(profiler)
+        self.assertFalse(result.success)
+        self.assertIn("MTL_CAPTURE_ENABLED", result.message)
+
+
+@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
+class TestSchedulerProfilerManagerMPS(unittest.TestCase):
+    """SchedulerProfilerManager._start_profile handles Metal capture failures."""
+
+    def _make_manager(self, output_dir):
+        from sglang.srt.managers.scheduler_components.profiler_manager import (
+            SchedulerProfilerManager,
+        )
+
+        class FakePS:
+            tp_rank = dp_rank = pp_rank = moe_ep_rank = 0
+            dp_size = pp_size = moe_ep_size = 1
+            gpu_id = 0
+
+        mgr = SchedulerProfilerManager(
+            ps=FakePS(), dp_tp_cpu_group=None, get_forward_ct=lambda: 0
+        )
+        mgr._init_profile(output_dir, None, None, None, None, None, False, "test")
+        return mgr
+
+    def test_start_profile_failure_does_not_crash(self):
+        import mlx.core as mx
+
+        from sglang.srt.hardware_backend.mlx.profiler import (
+            apply_metal_profiler_patches,
+        )
+
+        apply_metal_profiler_patches()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            mgr = self._make_manager(tmp)
+            with patch.object(
+                mx.metal,
+                "start_capture",
+                side_effect=RuntimeError("Capture layer is not inserted"),
+            ):
+                result = mgr._start_profile()
+
+        self.assertFalse(result.success)
+        self.assertFalse(mgr.profile_in_progress)
+        self.assertIsNone(mgr.torch_profiler)
+
+    def test_start_profile_success_with_mock_capture(self):
+        from unittest.mock import patch as mock_patch
+
+        import mlx.core as mx
+
+        from sglang.srt.hardware_backend.mlx.profiler import (
+            apply_metal_profiler_patches,
+        )
+
+        apply_metal_profiler_patches()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            mgr = self._make_manager(tmp)
+            with mock_patch.object(mx.metal, "start_capture"), mock_patch.object(
+                mx.metal, "stop_capture"
+            ), mock_patch("torch.distributed.barrier"):
+                result = mgr._start_profile()
+                self.assertTrue(result.success)
+                self.assertTrue(mgr.profile_in_progress)
+                mgr._stop_profile()
+                self.assertFalse(mgr.profile_in_progress)
+
+
+if __name__ == "__main__":
+    unittest.main()