From 0a28a929dccad374995220cb5d9f43ce77517e78 Mon Sep 17 00:00:00 2001 From: Lijuan Tang Date: Wed, 17 Jun 2026 13:06:20 -0700 Subject: [PATCH] [MLX] Add Metal profiling hooks to server profiler (#28122) --- .../srt/hardware_backend/mlx/profiler.py | 261 ++++++++++++++++++ .../scheduler_components/profiler_manager.py | 17 +- .../mlx/test_metal_profiler.py | 251 +++++++++++++++++ 3 files changed, 524 insertions(+), 5 deletions(-) create mode 100644 python/sglang/srt/hardware_backend/mlx/profiler.py create mode 100644 test/registered/unit/hardware_backend/mlx/test_metal_profiler.py diff --git a/python/sglang/srt/hardware_backend/mlx/profiler.py b/python/sglang/srt/hardware_backend/mlx/profiler.py new file mode 100644 index 000000000..c009ee149 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 31df519f9..42b47b598 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/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: diff --git a/test/registered/unit/hardware_backend/mlx/test_metal_profiler.py b/test/registered/unit/hardware_backend/mlx/test_metal_profiler.py new file mode 100644 index 000000000..40ff273fb --- /dev/null +++ b/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()