56 lines (37 loc) · 1.34 KB

ThunderMittens

ThunderMittens is an Apple Metal Shading Language (MSL) port for the ThunderKittens framework.

Project Structure

The repository supports two primary use cases:

1. MSL Kernel Development

For writing Metal Shading Language (MSL) kernels:

Clone the ThunderKittens repository
Open the project in Xcode
Xcode will handle all build processes

2. MLX Kernel Integration with Python

For using ThunderKittens kernels within MLX in Python:

Prerequisites

Python 3.8+
CMake
Xcode Command Line Tools

Installation Steps

Navigate to ThunderMittens/mlx directory

Install MLX with parallel build:

CMAKE_BUILD_PARALLEL_LEVEL=8 pip install -e ".[dev]"

Navigate to ThunderMittens/kernels directory
Install requirements:
```
pip install -r requirements.txt
```
Build kernels and bindings:
```
python setup.py build_ext -j8 --inplace
```

References

Please see our blog post to learn more about this work. Please checkout our paper to learn more about the ThunderKittens project.