Skip to content

mindtorch_v2/mps: real Metal GPU compute backend for 30+ ops#2755

Merged
lvyufeng merged 1 commit into
candle-org:masterfrom
lvyufeng:metal-gpu-backend
Mar 8, 2026
Merged

mindtorch_v2/mps: real Metal GPU compute backend for 30+ ops#2755
lvyufeng merged 1 commit into
candle-org:masterfrom
lvyufeng:metal-gpu-backend

Conversation

@lvyufeng

@lvyufeng lvyufeng commented Mar 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Replace NumPy CPU fallbacks with actual Metal GPU compute kernels dispatched via MTLComputeCommandEncoder
  • Add 105 MSL (Metal Shading Language) kernels covering element-wise binary (9), unary (16+2 activations), reduction (6), softmax, fill, copy, and in-place (8) operations — all with f32/f16 variants
  • Add MetalKernelDispatcher engine: lazy shader compilation, pipeline state caching, and thread dispatch helpers (PyObjC + ctypes dual-path, no new dependencies)
  • Implement GPU matmul via Apple's MPSMatrixMultiplication kernel — 4.3x faster than Accelerate BLAS at 1024×1024
  • All f64 and non-contiguous tensors automatically fall back to existing CPU/Accelerate path

New files

File Purpose
metal_shaders.py All MSL kernel source strings (generated from templates)
metal_compute.py GPU dispatch engine: compilation, pipeline cache, encode+commit

Modified files

File Changes
runtime.py +compile_library, make_compute_pipeline, get_compute_encoder (PyObjC + ctypes)
mps_kernels.py GPU matmul via MPSMatrixMultiplication (PyObjC + ctypes fallback)
ops.py 30+ ops upgraded: GPU-first dispatch with CPU fallback

Test plan

  • All 27 correctness checks pass (binary, unary, reduction, matmul, activations, in-place, softmax)
  • torch.allclose validation against CPU reference for every op
  • MatMul 1024×1024 benchmark: GPU 3.71ms vs CPU 16.10ms (4.3x speedup)
  • f64 tensors correctly fall back to CPU path
  • Non-contiguous tensors correctly fall back to CPU path
  • Run full model test suite (BERT, GPT-2) on MPS device

🤖 Generated with Claude Code

Replace NumPy CPU fallbacks with actual Metal GPU compute kernels
for element-wise, reduction, matmul, activation, and in-place ops.

- Add metal_shaders.py: 105 MSL kernels (f32/f16 variants)
- Add metal_compute.py: GPU dispatch engine with lazy compilation
  and pipeline caching (PyObjC + ctypes dual-path)
- Extend runtime.py: compile_library, make_compute_pipeline,
  get_compute_encoder with ctypes fallbacks
- Update mps_kernels.py: MPSMatrixMultiplication GPU matmul
  (4.3x faster than Accelerate BLAS at 1024x1024)
- Update ops.py: GPU dispatch for 9 binary, 16 unary, 4 reduction,
  4 NN fused, and 8 in-place ops with automatic f64/non-contiguous
  CPU fallback

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@lvyufeng
lvyufeng force-pushed the metal-gpu-backend branch from 251f208 to 903f61f Compare March 8, 2026 07:06
@lvyufeng
lvyufeng merged commit ec2520e into candle-org:master Mar 8, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant