Add Apple MPS backend #1
Workflow file for this run
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Build MPS Backend | |
| on: | |
| push: | |
| branches: [apple-mps-backend] | |
| pull_request: | |
| branches: [main] | |
| paths: | |
| - 'third_party/apple/**' | |
| - 'CMakeLists.txt' | |
| - 'setup.py' | |
| jobs: | |
| build-mps: | |
| runs-on: macos-26 | |
| timeout-minutes: 60 | |
| env: | |
| TRITON_BUILD_WITH_CLANG_LLD: "FALSE" | |
| steps: | |
| - name: Checkout | |
| uses: actions/checkout@v4 | |
| with: | |
| submodules: recursive | |
| - name: Install Metal Toolchain | |
| run: | | |
| # -downloadComponent is Xcode 26+ only; on older Xcode, Metal tools are bundled | |
| xcodebuild -downloadComponent MetalToolchain || echo "Metal toolchain already bundled" | |
| xcrun --find metal-as || echo "WARNING: metal-as not found" | |
| - name: Install brew dependencies | |
| run: | | |
| brew install ccache coreutils ninja | |
| - name: Build MetalASM | |
| run: | | |
| cd third_party/apple/MetalASM | |
| swift build -c release --product MetalASMBridge | |
| ls -la .build/release/libMetalASMBridge.dylib | |
| - name: Create venv and install deps | |
| run: | | |
| python3 -m venv ~/.venv | |
| source ~/.venv/bin/activate | |
| pip install --upgrade pip | |
| pip install ninja cmake wheel numpy pytest pybind11 setuptools | |
| # Install PyTorch with PrecompiledMetalShaderLibrary API | |
| pip install https://github.com/imperatormk/pytorch/releases/download/v0.1.0-mps-metallib/torch-2.12.0a0+git440ad94-cp314-cp314-macosx_26_0_arm64.whl | |
| - name: Install Triton (Apple MPS only) | |
| env: | |
| TRITON_BUILD_WITH_O1: "true" | |
| MAX_JOBS: 3 | |
| NINJA_STATUS: "[%f/%t, %es elapsed] " | |
| run: | | |
| source ~/.venv/bin/activate | |
| # Let Triton download and build its pinned LLVM | |
| pip install -e . --no-build-isolation -v | |
| - name: Verify import | |
| run: | | |
| source ~/.venv/bin/activate | |
| python -c " | |
| import triton | |
| from triton.backends.apple.compiler import MPSBackend | |
| print('Triton MPS backend imported successfully') | |
| print(f'Triton version: {triton.__version__}') | |
| " |