执行摘要

本PR为sglang仓库添加了原生MLX执行后端，专为Apple Silicon Mac设计，通过环境变量SGLANG_USE_MLX=1激活。核心变更包括引入MlxModelRunner和MlxTpModelWorker，实现端到端MLX推理，避免PyTorch MPS开销，显著提升性能。此功能已合并，影响范围限于Mac用户，但展示了硬件后端集成的优雅模式。

功能与动机

动机是提升在Apple Silicon Mac上的推理性能。PR body中明确表示：“Introduces MlxModelRunner and MlxTpModelWorker under python/sglang/srt/hardware_backend/mlx, enabling end-to-end model inference via MLX on Apple Silicon.” 此外，Issue评论中提供了性能数据，显示使用MLX后吞吐量改善，验证了需求。

实现拆解

实现按模块拆解：

• 硬件后端模块：新增python/sglang/srt/hardware_backend/mlx/model_runner.py和tp_worker.py，分别负责MLX推理和调度集成。
• 工具模块：新增python/sglang/srt/utils/tensor_bridge.py，提供零拷贝张量转换，关键函数如torch_to_mlx和mlx_to_torch。
• 调度模块：修改python/sglang/srt/managers/scheduler.py，在init_tp_model_worker中根据use_mlx()选择工作者。
• 基准测试：修改python/sglang/bench_one_batch.py，引入_BenchRunner抽象，统一PyTorch和MLX路径。
• 其他修改：如python/sglang/_mps_stub.py添加StreamContext，python/sglang/jit_kernel/diffusion/triton/mps_fallback.py集成MLX加速norm函数。

评论区精华

Review讨论中亮点包括：

• 内存效率：gemini-code-assist[bot]指出初始实现可能加载模型两次，作者通过添加MlxModelRunnerStub解决，跳过PyTorch权重加载。
• 正确性：gemini-code-assist[bot]提到MLX后端未支持ForwardMode.MIXED，可能导致错误；此问题未解决，需后续处理。
• 性能担忧：alexnails询问OOM行为，作者打开issue #21443跟踪，显示内存管理风险。
• 代码风格：讨论中涉及benchmark命名和类型提示修正，体现对细节的关注。

风险与影响

风险：

• 兼容性：仅限Apple Silicon，其他平台无影响；但张量桥接中的dtype映射可能不完整。
• 性能：MLX后端在不同场景下可能表现不一致；benchmark序列处理可能不准确。
• 回归：修改调度器可能影响现有MPS后端；缺少全面测试覆盖。

影响：

• 用户：Mac用户获得性能提升，通过简单环境变量启用。
• 系统：代码库增加新后端，但设计最小化核心变更。
• 团队：需维护MLX代码，CI已添加相关标签确保测试。

关联脉络

与此PR相关的历史PR包括：

• 20782：添加StreamContext stub，同为MPS支持，显示对Apple Silicon的持续优化。

• 20753：支持sglang.check_env，增强Mac兼容性。

• 18032：为NPU添加Hybrid KV Cache，类似硬件后端扩展，设计模式可借鉴。

  
  这些PR共同展示了sglang向多硬件平台演进的趋势。