skmetal

Apple Silicon GPU acceleration for scikit-learn

Decorate any scikit-learn estimator with @skmetal.accelerate and fit()/predict() run on the GPU — no code changes. Leverages Apple Silicon's unified memory for zero-copy data sharing between numpy and Metal.

import skmetal
from sklearn.linear_model import LinearRegression

@skmetal.accelerate
def model():
    return LinearRegression()

m = model()
m.fit(X_train, y_train)
m.predict(X_test)

Installation

pip install skmetal

macOS 14+ and Apple Silicon (M1–M5) required. No Xcode needed — the pip package includes a pre-built dylib and Metal library.

From source (for development)

git clone https://github.com/abderahmane-ai/skmetal.git
cd skmetal
pip install -e ".[dev]"

# Build Swift + Metal (required after any .metal or .swift change)
cd skmetal_bridge
bash compile_metal.sh
swift build --configuration release
cp .build/arm64-apple-macosx/release/libSkMetalBridge.dylib ../skmetal/
cd ..

Benchmarks

Measured on an M4 Air (16 GB) and M4 Max (128 GB). All data float32, n_init=1, max_iter=30 for KMeans.

KMeans MLX requires pip install skmetal[mlx]. Without MLX, KMeans falls back to a Metal fused command-buffer (slower than CPU — 0.1×). The MLX path uses flash-kmeans-mlx which fuses distance + argmin + update into a single compiled GPU kernel.

Run benchmarks locally:

pip install skmetal[mlx]
python benchmarks/run_compare.py     # moderate data sizes
python benchmarks/benchmark_suite.py # large data (generates baseline.json)

Features

• Zero-copy GPU execution — numpy arrays passed directly to Metal via bytesNoCopy on unified memory
• Drop-in acceleration — decorate any estimator-returning function, wrap an existing instance, or use a context manager
• Smart dispatch — automatically routes to CPU for small datasets where GPU overhead dominates; configurable per-estimator thresholds
• GPU solvers — Cholesky, FISTA, L-BFGS, IRLS, KNN tile-then-merge, SIMD-group GEMM, flash-kmeans-mlx compiled kernels
• float4 vectorization — 6 kernel families use float4 loads/stores for 4× memory throughput
• Optional MLX acceleration — install skmetal[mlx] for flash-kmeans-mlx GPU KMeans (3–30×) and TruncatedSVD (GPU SVD)
• Transparent fallback — imports cleanly on non-Apple-Silicon machines; all operations fall back to scikit-learn CPU
• Verbose logging — skmetal.set_verbose(True) prints why each estimator chose GPU or CPU

Usage

Decorator (recommended)

Wrap any function that returns an estimator:

import skmetal
from sklearn.linear_model import LogisticRegression

@skmetal.accelerate
def model():
    return LogisticRegression(random_state=42)

clf = model()
clf.fit(X_train, y_train)
clf.predict(X_test)

Works with pipelines too:

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

@skmetal.accelerate
def pipe():
    return Pipeline([
        ("scaler", StandardScaler()),
        ("clf", LogisticRegression()),
    ])

p = pipe()
p.fit(X, y)

Function call

model = skmetal.accelerate(LinearRegression())
model.fit(X, y)

Context manager

with skmetal.accelerate_context():
    model = LinearRegression()
    model.fit(X, y)

Checking GPU availability

import skmetal

if skmetal.METAL_AVAILABLE:
    info = skmetal.device_info()
    print(info)
    # {'name': 'Apple M4 Max', ...,
    #  'has_unified_memory': True,
    #  'recommended_working_set_size_bytes': 68719476736}

Configuration

import skmetal

skmetal.set_device("cpu")             # force CPU fallback globally
skmetal.set_verbose(True)             # log dispatch decisions
skmetal.set_threshold(100_000)        # global min n*d for GPU
skmetal.update_threshold("KMeans",    # per-estimator override
                         min_rows=100_000, min_cols=50)
skmetal.reset_thresholds()            # restore defaults

config = skmetal.get_config()
print(config)

On non-Apple-Silicon machines skmetal imports cleanly and all estimators transparently fall back to scikit-learn CPU.

Supported Estimators

Each estimator has per-estimator (min_rows, min_cols) thresholds. Below the threshold the estimator uses CPU. Override via skmetal.update_threshold() or force GPU with skmetal.set_device("gpu").

Architecture

numpy array → np.ctypes.data → UnsafeMutableRawPointer → MTLBuffer(bytesNoCopy:) → Metal GPU
                   |                                                    |
                   +--------- same physical memory (unified) -----------+

Apple Silicon's unified memory enables zero-copy data sharing. The Swift bridge exposes @_cdecl functions callable from Python via ctypes. Each estimator's fit()/predict() calls the appropriate bridge function, which dispatches Metal Performance Shaders or custom compute kernels.

Metal kernels (13 files)

Swift bridge (9 files)

Project Structure

skmetal/
  skmetal/
    __init__.py          # public API: accelerate, config, device_info
    _about.py            # version
    _bridge.py           # ctypes → Swift @_cdecl exports (47 functions)
    _config.py           # Config dataclass, thresholds, device control
    _dispatch.py         # estimator registry + wrapping logic
    accelerate.py        # @accelerate decorator + context manager
    estimators/
      _base.py           # BaseGPUEstimator with fallback logic
      _registry.py       # GPU_REGISTRY (19 estimators)
      _mlx_registry.py   # MLX detection
      _mlx_svd.py        # TruncatedSVD MLX backend
      linear_model.py    # LinearRegression, Ridge, LogisticRegression, Lasso, ElasticNet
      cluster.py         # KMeans, DBSCAN
      decomposition.py   # TruncatedSVD
      ensemble.py        # HistGradientBoosting
      naive_bayes.py     # GaussianNB
      neighbors.py       # KNeighbors, NearestNeighbors
      preprocessing.py   # StandardScaler, MinMaxScaler, RobustScaler
      svm.py             # SVC, SVR
  skmetal_bridge/        # Swift + Metal SPM package
    Sources/SkMetalBridge/
      Bridge.swift
      LinearModelBridge.swift
      KMeansBridge.swift
      KNNBridge.swift
      LinearAlgebraBridge.swift
      PreprocessingBridge.swift
      MinMaxBridge.swift
      LogisticBridge.swift
      SVCBridge.swift
      SVTreeBridge.swift
      MetalContext.swift
      Kernels/*.metal    # 13 Metal kernel files
  benchmarks/
    run_compare.py       # quick comparison benchmark
    benchmark_suite.py   # full benchmark suite (generates baseline)
    baseline.json
  tests/                 # 263 tests across 11 files
  pyproject.toml
  LICENSE
  .github/workflows/
    ci.yml               # build + ruff + pytest + benchmarks
    release.yml          # PyPI publish on v* tag

Development

git clone https://github.com/abderahmane-ai/skmetal.git
cd skmetal

# Build Swift + Metal (required after any .metal or .swift change)
cd skmetal_bridge
bash compile_metal.sh
swift build --configuration release
cp .build/arm64-apple-macosx/release/libSkMetalBridge.dylib ../skmetal/
cd ..

# Install in editable mode
pip install -e ".[dev]"

# Lint
ruff check skmetal/ tests/
ruff format --check skmetal/ tests/

# Run tests (from skmetal/ dir)
cd skmetal && python3 -m pytest ../tests/ -q --tb=short

# Run benchmarks
python benchmarks/run_compare.py

Adding a new estimator

1. Create skmetal/estimators/my_model.py with MetalMyModel(BaseGPUEstimator)
2. Implement fit()/predict() calling the Swift bridge via _bridge_call()
3. Register in estimators/_registry.py → GPU_REGISTRY
4. Add a parametrized test in tests/test_correctness.py
5. Write a Metal kernel + Swift bridge function if needed

License

MIT