Escape from ABI Hell: Ditching C++ and Nanobind for Pure-Python Metal JIT

If you set out to write a custom GPU kernel for MLX, everyone gives you the “sensible adult” advice.

You scaffold a C++ project. You pull in a binding library like nanobind or pybind11. You write a CMakeLists.txt file that tries to locate your local Metal compiler (xcrun metal). And you write boilerplate wrappers to cast Python tensor handles into native C++ references before dispatching to the hardware.

It is the classic “if it ain’t broke, don’t fix it” approach. Nobody ever got fired for following the established pattern.

Except sometimes, if you actually want to extract every drop of performance out of unified memory, you have to break some eggs.

The “Sensible” Path Broke Me

When I started building mlx-ternary, I followed the textbook playbook. The mission was straightforward: execute custom ternary matrix multiplication (W ∈ {-1, 0, 1} across M × K) against FP16 activation vectors.

The initial setup looked respectable:

TernaryOp.cpp: Defined custom MLX C++ primitives and bound them via nanobind.
ternary.metal: Contained the raw Metal compute shaders.
CMakeLists.txt: Orchestrated compilation into a .so extension.

The local debug build compiled cleanly. The test suite passed. I felt like a responsible systems engineer.

Wait. This felt too easy.

That was me staring at my terminal.

Because getting a build to pass on your own laptop wasn’t the goal. The guiding principle was hardware optimization: squeeze maximum bandwidth out of Apple’s unified memory architecture. The ends had to justify the means.

Instead of optimization, I got an obstacle course.

1. Welcome to ABI Hell (Or: Why My Terminal Exploded)

Before getting into the crash, what does ABI actually mean?

We talk about APIs all day. An API (Application Programming Interface) is the polite agreement in source code: “Here is a function named gemv, it takes three tensor arrays, and returns an output array.” If your syntax matches the function signature, your compiler gives you a thumbs up.

The ABI (Application Binary Interface) is what happens when the compiler finishes being polite.

It is the raw, physical memory layout in machine code. How many bytes wide is that tensor struct? Where does the pointer to the underlying data buffer sit in RAM? Which ARM64 registers hold the arguments? How does the compiler mangle function names into binary linker symbols?

In pure Python, you never think about this because the interpreter handles memory objects. In pure C++, the compiler links everything together in a single binary.

The nightmare begins when you build a C++ Python extension. Now, you have two separately compiled binary worlds trying to share the exact same memory space:

_mlx.so: The precompiled MLX wheel built by Apple on their CI machines.
ternary.so: Your custom extension compiled on your laptop with your local Clang.

When you pass an mlx.core.array from Python into your C++ function, your extension has to verify the object’s type via C++ Run-Time Type Information (RTTI).

If your local Clang compiler version, C++ standard library (libc++), or nanobind headers diverge by even a minor point release from what Apple used to build MLX: the memory layouts do not match. The dynamic linker looks at the binary symbols across the shared library boundary, gets confused, and points the CPU at the wrong byte address.

You don’t get an informative Python traceback with a line number. You get an immediate, uncatchable SIGSEGV or std::bad_cast core dump. The process just dies.

That is ABI hell. Your C++ code is valid. Your tensor math is correct. But the two compiled binaries are speaking slightly different dialects of C++ in RAM, and a little Faiz cannot keep up with Apple’s internal release velocity.

Dependency hell had officially become my own personal hell. I spent two entire days tracking down memory faults, only to discover nanobind had to be pinned strictly to v2.10.2 to match MLX’s internal type registry.

2. The Xcode Homework Tax

As if debugging tensor math wasn’t hard enough.

This project was supposed to be learning by building. Instead, I found myself pulling my hair out in terminal buffers. Apple treats developer environments like moving targets. A minor macOS point release silently shifts SDK headers in /Library/Developer/CommandLineTools.

If an engineer clones your repository to test benchmarks on their M3 Max, they are instantly greeted by 80 lines of CMake red text:

CMake Error at /usr/local/share/cmake/Modules/FindPackageHandleStandardArgs.cmake:
  Could NOT find Metal (missing: METAL_LIBRARY METAL_INCLUDE_DIR)

Now you are not debugging tensor math. You are debugging why another developer’s Xcode command line tools have an invalid symlink to the macOS 15.1 SDK.

This completely defeats the entire pitch of local Apple Silicon AI. The promise was: “Hey, look how fast you can run models on a consumer laptop to close the gap with Nvidia machines!”

The reality was: “Here is your 12 GB Xcode homework. Agree to this 80-page license agreement, reboot twice, and pray your Clang version matches my wheel.”

If your lightweight local tool requires half an hour of Xcode triage before token zero, your tool is dead on arrival.

(Geez, Faiz, so dramatic. But seriously…)

Bypassing the FFI Layer: Discovering `metal_kernel`

That was when the realization hit: We did not need C++ at all.

(Disclaimer: This is not a generalization for all software ever built, so please do not come @ me on social media. C++ is great. C++ pays mortgages.)

But in this specific case, why were we compiling native C++ binaries? Only to serve as a glorified middleman. The C++ layer was doing nothing more than taking a pointer to a tensor and handing it to Apple’s Metal driver.

Apple Silicon has a unique physical reality: the GPU and CPU share the same unified RAM pool. There is no PCIe bus transfer. There is no host-to-device memory staging. A buffer allocated in Python via MLX already lives directly in memory that the Apple GPU can read.

Buried inside MLX’s low-level API sits an alternative: mlx.core.fast.metal_kernel.

Instead of compiling C++ ahead of time, fast.metal_kernel lets you compile raw Metal Shading Language (MSL) source code directly from Python at runtime.

Just in time. Literally.

import mlx.core as mx
import mlx.core.fast as fast

# Compile raw Metal C++ directly from Python at runtime
kernel = fast.metal_kernel(
    name="ternary_gemv_lut",
    input_names=["activations", "weights", "scales"],
    output_names=["output"],
    source=get_gemv_source(M, K, N, group_size),
    header=_LUT_HEADER
)

The first time this executes, MLX hands the MSL string directly to the macOS Metal compiler service. The driver compiles it into machine instructions tailored for your specific GPU architecture and caches the pipeline state object in memory.

Subsequent invocations cost single-digit microseconds. No C++ compiler setup. No CMake build trees. No nanobind version pinning.

(And for anyone who felt a little lost while reading through the linker jargon: deep breath, no more pain. I am genuinely trying to make these writeups accessible for the four people who read them :))

How the Pure-Python Metal Pipeline Works

Here is the entire runtime execution lifecycle inside mlx_ternary/ops.py:

Rendering Graph...

1. In-Memory Kernel Specialization

When you compile ahead of time in C++, you either generate bloated binaries with dozens of template instantiations, or you accept dynamic branches inside your GPU shader.

With runtime JIT, we specialize the shader code to the exact matrix dimensions of the layer being decoded:

_KERNEL_CACHE = {}

def get_cached_gemv(M: int, K: int, N: int, group_size: int):
    key = (M, K, N, group_size)
    if key not in _KERNEL_CACHE:
        # Bake exact matrix dimensions directly into shader constants
        source = f"""
        uint tid = thread_position_in_grid.x;
        uint row = thread_position_in_grid.z;
        if (row >= {M} || tid >= {N}) return;

        // Unroll loops with zero runtime branching
        uint K_blocks = {K} / {group_size};
        ...
        """
        _KERNEL_CACHE[key] = fast.metal_kernel(
            name=f"gemv_{M}_{K}_{N}_{group_size}",
            input_names=["x", "w", "scales"],
            output_names=["out"],
            source=source,
            header=_LUT_HEADER
        )
    return _KERNEL_CACHE[key]

By substituting {M}, {K}, and {group_size} directly into the MSL string prior to compilation, the Metal compiler unrolls inner loops and hardcodes register limits. The GPU executes with zero dynamic branch overhead.

2. Direct Unified Memory Dispatch

Because the buffers live in unified memory, dispatching the kernel from Python requires no data copying:

def ternary_gemv(x: mx.array, w: mx.array, scales: mx.array) -> mx.array:
    M, K = x.shape
    N = w.shape[0]
    kernel = get_cached_gemv(M, K, N, group_size=128)
    
    # Configure thread grid matching Apple Silicon execution geometry
    grid = (N, 1, M)
    threadgroup = (min(N, 1024), 1, 1)
    
    return kernel(
        inputs=[x, w, scales],
        grid=grid,
        threadgroup=threadgroup,
        output_shapes=[(M, N)],
        output_dtypes=[x.dtype]
    )[0]

The call returns an mx.array ready for downstream layers. The Python overhead of the wrapper is under 4 microseconds. Woot woot.

The Payoff: Watching the Sunrise on a Grateful Universe

And now, we finally rest, watching the sunrise over a grateful universe.

The architectural pivot from C++/CMake to runtime MSL JIT brought three immediate practical wins:

True Wheel Portability: mlx-ternary is a standard Python package. The wheel contains pure .py files. It installs on any Apple Silicon Mac in 200 milliseconds without triggering local compilation.
Zero Toolchain Prerequisites: End users don’t need Xcode, CMake, or a local Clang toolchain installed. If macOS has Metal (standard on every Mac), the kernel runs.
Zero Latency Penalty: Removing the C++ glue layer cost nothing in raw compute. In fact, because the JIT compiler bakes matrix strides directly into MSL constants at runtime, our kernel executes in 365 µs, outperforming Apple’s native 2-bit quantized matmul (519 µs) by 1.42x on M4 Pro hardware.

(Apple, my DMs are open. Just saying.)

Compiling C++ extensions was an old reflex carried over from CUDA and Linux x86 ecosystems.

It wasn’t my habit. But coming at it as an outsider, avoiding the traditional FFI tax turned out to be the cleanest path to real gains.

Either that, or I got lucky and only time will tell.

In the next log, we examine The Horizon Architecture: how to pack five base-3 trits into an 8-bit byte (3⁵ = 243 ≤ 255) using a constant-memory register LUT to saturate 273 GB/s memory bandwidth.