Prhub

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

原始 PR 作者 LijuanTang94 合并时间 2026-06-18 04:06 文件变更 3 提交数 11 评论 18 代码增减 +524 / -5

执行摘要

新增 Apple Silicon Metal 性能分析后端

扩展 #22159 首次添加的 Apple Silicon profiling 支持至服务端 profiler,让 bench_serving.py 和 bench_one_batch_server.py 也能启用 Metal GPU trace 捕获,以完成 Apple Device Support 路线图(issue #19137)中 Profiling 部分的目标。

值得精读。该 PR 展示了如何通过 backend patch 模式在 SGLang 中为新型硬件添加 profiling 支持,设计思路清晰,与已有 NPU 后端一致,可作为跨后端扩展的参考模板。对于 Apple Silicon 开发人员,还提供了详细的用法和验证结果。

讨论亮点

review 中 yeahdongcn 提出了关键设计建议:

  • 使用 is_mps() 条件代替 use_mlx() 以同时支持 MPS 和 MLX 模式,避免遗漏纯 torch MPS 用户。
  • 将 Apple-specific profiling 逻辑从 common profiler 中拆出,采用与 NPU 类似的 backend patch 模式,降低主路径耦合。
  • 要求处理 start_mps/start_mlx 可能抛出的 RuntimeError,确保意外失败时不传播异常。
  • 要求补充单元测试,作者后新增 10 个测试。
  • 建议将 ProfileManager 导入移到文件顶部,作者已调整。所有建议均被采纳并实现,最终设计得到 approval。

实现拆解

  1. 新增 hardware_backend/mlx/profiler.py:实现 MetalCaptureProfiler 数据类,管理 MLX(mx.metal.start_capture)和 MPS(torch.mps.profiler.metal_capture)两条捕获路径的启动与停止;MetalTorchProfiler 作为 torch.profiler.profile 的替代包装,代理 start/stop 和 trace 导出;apply_metal_profiler_patches() 全局替换 torch.profiler.profile,并保证幂等性。
  2. 修改 profiler_manager.py:在模块加载时新增 _is_mps 判断,在 NPU 补丁之后调用 apply_metal_profiler_patches();将 ProfileManager 导入移至文件顶部;在 _start_profile 中将 self.torch_profiler.start() 包裹在 try/except RuntimeError 中,捕获失败时返回 success=False 而非抛出异常。
  3. 新增 test_metal_profiler.py:10 个单元测试覆盖补丁替换、幂等性、MLX/MPS 路径成功与 RuntimeError 降级、以及 SchedulerProfilerManager 的完整 start/stop 周期,所有测试自动跳过非 Apple Silicon 或缺少 mlx 的环境。
文件 模块 状态 重要度
python/sglang/srt/hardware_backend/mlx/profiler.py 硬件后端 added 9.25
test/registered/unit/hardware_backend/mlx/test_metal_profiler.py 单元测试 added 8.14
python/sglang/srt/managers/scheduler_components/profiler_manager.py 调度器 modified 6.68

关键符号

MetalCaptureProfiler.start_mlx MetalCaptureProfiler.start_mps MetalCaptureProfiler.stop MetalTorchProfiler.start MetalTorchProfiler.stop apply_metal_profiler_patches SchedulerProfilerManager._start_profile

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

MPS/MLX 后端架构选择 设计

yeahdongcn 建议将 Apple-specific profiling 从 common profiler 中分离,采用类似 NPU 的 backend patch 模式,并同时支持 MPS 和 MLX 两种模式。作者初始版本仅在 common profiler 中添加 use_mlx 分支,后续重构成独立的 hardware_backend/mlx/profiler.py 并通过 apply_metal_profiler_patches() 注入。

结论:采纳重构意见,采用 backend patch 模式,profiler_manager.py 仅添加 3 行 elif _is_mps 分支。 · 已解决

优雅处理 Metal 捕获失败 正确性

yeahdongcn 指出 start_mps/start_mlx 可能抛出 RuntimeError(如未设置 MTL_CAPTURE_ENABLED),若不处理会传播到调用者。作者在 MetalTorchProfiler.start() 和 profiler_manager._start_profile() 中均添加了 try/except,将异常转换为 ProfileReqOutput(success=False) 并给出可操作提示。

结论:异常被捕获并优雅降级,不会导致服务崩溃。 · 已解决

添加单元测试 测试

yeahdongcn 要求为 profiler.py 添加单元测试。作者随后新增 test_metal_profiler.py,包含 10 个测试,覆盖补丁替换、幂等性、MLX/MPS 路径成功与失败、以及完整 start/stop 周期。

结论:测试已添加并通过。 · 已解决

导入顺序优化 style

yeahdongcn 指出 ProfileManager 导入位置应移至文件顶部以避免模块级 import 混乱。作者将导入从末尾移到与其他 import 一起。

结论:已调整。 · 已解决

风险与影响

风险较低。主要风险包括:

  • 新增 Metal 捕获路径完全条件化(_is_mps),不影响现有 CUDA/XPU/NPU 后端。
  • 在非 Apple Silicon 平台上 is_mps() 为 False,补丁不会被应用。
  • Metal 捕获失败(如未设置 MTL_CAPTURE_ENABLED)时通过 try/except 优雅降级返回 success=False,不会导致服务崩溃。
  • MPS 路径依赖 torch.mps.profiler.metal_capture(需 torch >= 2.1),代码中有 hasattr 检查,兼容性较好。
  • 测试覆盖了 MLX 和 MPS 的 mock 路径,但未在真实硬件上集成测试,建议在 CI 中添加 Apple Silicon runner 以验证。

对 Apple Silicon 用户是纯增益:可通过 SGLANG_TORCH_PROFILER_DIR 环境变量和 --profile 参数在服务端 profiling 时获得 Metal GPU trace(.gputrace 文件)。两种模式均可:MLX 模式(SGLANG_USE_MLX=1)使用 mx.metal API,MPS 模式(默认)使用 torch.mps.profiler API。生成 trace 需设置 MTL_CAPTURE_ENABLED=1 环境变量。不影响非 Apple Silicon 平台,性能影响仅限于 profiling 开启时的额外开销。代码设计易于后续新增其他硬件后端。

新硬件后端 依赖 MLX 包 异常处理优雅降级

关联 Issue

#19137 [Roadmap] Apple Device Support (2026 Q1)

完整报告

参与讨论