Skip to content
 
 

Repository files navigation

QuixiCore Metal

QuixiCore Metal is the Apple Silicon backend for the QuixiCore kernel library. It provides native Metal kernels with Python integrations for MLX and PyTorch MPS.

The backend follows the shared QuixiCore contract: common operation names, correctness expectations, quantization metadata, and benchmark conventions, implemented natively for Apple GPUs.

What Is Included

  • Metal Shading Language kernels under kernels/ and include/metal/.
  • MLX Python bindings exposed as tk.
  • PyTorch MPS bindings exposed as tk_torch.
  • Xcode project support through QuixiCoreMetal.xcodeproj.
  • Correctness, parity, and benchmark harnesses for the supported integrations.

Kernel coverage includes normalization, activation, attention, linear attention, state-space, matmul, quantization, vision, MoE, sampling, serving, optimizer, and utility operations. The exact supported surface is tracked in .quixicore/kernels.yaml.

Requirements

  • Apple Silicon Mac.
  • Xcode with the Metal Toolchain installed.
  • Python virtual environment for Python integrations.

One-time Xcode and Metal setup

Install the full Xcode application, then point the command-line developer tools at it and complete Xcode's first-launch setup:

sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer
sudo xcodebuild -runFirstLaunch

If Xcode is installed somewhere other than /Applications/Xcode.app, replace that path with the location of your Xcode application. Verify the active developer directory and tools:

xcode-select -p
xcodebuild -version
xcrun --find metal

The active developer directory should be /Applications/Xcode.app/Contents/Developer, not /Library/Developer/CommandLineTools. Install the separately distributed Metal Toolchain component before the first build:

xcodebuild -downloadComponent MetalToolchain

xcrun --find metal may report a launcher path even when this component is not installed. If a build reports cannot execute tool 'metal' due to missing Metal Toolchain, run the download command above and retry.

To use full Xcode for only the current command instead of changing the system selection, set DEVELOPER_DIR:

DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer scripts/configure
DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer \
  scripts/build xcode -configuration Debug

The MLX binding currently targets the MLX 0.21 C++ extension API. Use Python 3.12 for the MLX path unless you are intentionally porting the extension to a newer MLX C++ API.

Build

Run commands from the repository root.

scripts/configure
scripts/build xcode -configuration Debug

For the MLX-backed Python package:

python3.12 -m venv .venv
. .venv/bin/activate
python -m pip install -r bindings/python/requirements.txt
PYTHON=.venv/bin/python scripts/build python

For the PyTorch MPS package:

python3 -m venv .venv-torch
. .venv-torch/bin/activate
python -m pip install torch
PYTHON=.venv-torch/bin/python scripts/build pytorch_mps

The PyTorch package builds its Objective-C++ extension and Metal library on first import.

Use

MLX:

import mlx.core as mx
import tk

x = mx.random.normal((4096, 1024)).astype(mx.bfloat16)
w = mx.ones((1024,), dtype=mx.bfloat16)
y = tk.rms_norm(x, w)
mx.eval(y)

PyTorch MPS:

import torch
import tk_torch

x = torch.randn(2, 128, 1024, dtype=torch.bfloat16, device="mps")
w = torch.ones(1024, dtype=torch.bfloat16, device="mps")
b = torch.zeros(1024, dtype=torch.bfloat16, device="mps")
y = tk_torch.layernorm(x, w, b)
torch.mps.synchronize()

Specialized composed operations

The public tk module also exposes pure tensor operations that combine common decode, sparse-projection, embedding, and vision stages without introducing an application-level runtime:

Operation Public API and contract
Packed embeddings quantized_embedding and quantized_embedding_bag gather or reduce GGUF/MX/FP rows directly from packed tables.
Decode projections decode_linear_epilogue and decode_swiglu support dense, q4_0, q8_0, q6_K, MXFP8, NVFP4, and MXFP4 weights, fused activations/bias/residuals, and optional output quantization.
Fused output sampling lm_head_sample supports q4_0/q8_0/MXFP8/NVFP4/MXFP4 projection with argmax, categorical, top-k, and top-p modes without materializing logits; q6_K supports fused argmax and categorical modes.
Sparse output projection lm_head_masked consumes packed allow masks; lm_head_candidates consumes CSR candidate lists. Both support dense, q4_0, q8_0, q6_K, MXFP8, NVFP4, and MXFP4 weights and return deterministic top-k ids and log-probabilities without materializing full logits.
Quantized beam advance lm_head_beam_advance combines q4_0/q8_0/MXFP8/NVFP4/MXFP4 output projection, exact full-vocabulary normalization, cumulative beam scores, and deterministic parent/token selection.
Spatial projection space_to_depth_norm_linear composes block-2/block-4 space-to-depth, LayerNorm, and projection with odd-edge padding.
Pairwise edge MLP edge_mlp_256x7 factorizes a fixed 512→256→7 pairwise MLP so the first projection scales with sequence length rather than pair count.
Head-major GQA decode attn_decode_bh consumes preallocated (B,Hkv,cache_T,D) caches and partitions long contexts across SIMD groups.
Functional cache decode decode_cache_attention composes optional Q/K RMSNorm, split-half RoPE, functional cache append, and GQA attention.

MLX arrays and PyTorch MPS tensors use the same top-level functions. Operations with measured crossover points auto-route between direct Metal and framework composition; use_kernel=True or False selects a path explicitly where the API exposes that option.

Test

# Xcode build-for-testing
scripts/test xcode

# MLX correctness
PYTHON=.venv/bin/python scripts/test correctness

# PyTorch MPS correctness
PYTHON=.venv-torch/bin/python scripts/test mps

# Cross-backend parity; install torch in the MLX venv first if needed.
PYTHON=.venv/bin/python scripts/test parity

MPS and parity targets exit cleanly with a skip message when Torch or MPS support is unavailable.

Benchmark

Use perf/bench_kernels.py from the repository root:

PYTHON=.venv/bin/python perf/bench_kernels.py --backend mlx --preset smoke --kernel all
PYTHON=.venv-torch/bin/python perf/bench_kernels.py --backend torch --preset smoke --kernel all

Benchmark results are hardware-, OS-, framework-, and shape-dependent. See perf/perf.md for methodology and benchmark conventions.

Repository Layout

.quixicore/              Backend and kernel metadata
include/metal/           Shared Metal tile substrate and headers
include/quixicore/metal/ Public QuixiCore Metal headers
kernels/                 Operation implementations by family
bindings/python/         MLX-backed Python package
bindings/pytorch_mps/    PyTorch MPS package
bindings/mlx/            MLX source integration
tests/                   Correctness, parity, integration, and unit tests
perf/                    Benchmark harnesses, configs, results, and baselines
scripts/                 Common build, test, bench, and clean entry points

More detail is in docs/repository-structure.md.

Metadata And Docs

Credits

QuixiCore Metal builds on the ThunderKittens-style tiled GPU programming model originally published by HazyResearch and adapted for Apple Metal by QuixiAI. Some model-serving kernels were informed by reference implementations from metal-forge by AlpinDale.

License

MIT. See LICENSE.

About

No description, website, or topics provided.

Resources

Contributing

Security policy

Stars

17 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages