mlxcel-core: MLX C++ FFI Bridge¶

Table of Contents¶

• Why Direct C++ Bindings?
• cxx Bridge Architecture
• Module Overview
• Key Types
• Core Modules
• Build System
• Adding New Operations

Why Direct C++ Bindings?¶

MLX is an Apple Silicon-native ML framework written in C++. It already has a C wrapper layer (mlx-c) that provides a stable ABI, but this introduces overhead:

1. Allocation: the C wrapper allocates opaque handle objects on the heap for every array reference.
2. Copies: some operations require intermediate copies to cross the C ABI boundary.
3. Missing operations: the C wrapper does not expose every MLX operation, requiring workarounds.

mlxcel bypasses the C wrapper entirely and binds directly to the MLX C++ API via the cxx crate. This provides:

• Zero overhead for array passing (C++ UniquePtr<MlxArray> maps directly to the underlying mlx::core::array).
• Access to the full MLX operation set.
• The ability to write custom MLX kernels and register them from Rust.

The mlxcel-core crate lives at src/lib/mlxcel-core/.

cxx Bridge Architecture¶

The bridge is defined in src/lib/mlxcel-core/src/lib.rs using the cxx crate's #[cxx::bridge] macro:

#[cxx::bridge(namespace = "mlx_cxx")]
mod ffi {
    unsafe extern "C++" {
        include!("mlx_cxx_bridge.h");

        /// Opaque wrapper for mlx::core::array
        type MlxArray;

        /// Opaque wrapper for mlx::core::Stream
        type MlxStream;

        // Operations are declared here and implemented in C++
        fn matmul(a: &MlxArray, b: &MlxArray) -> UniquePtr<MlxArray>;
        // ...
    }
}

The mlx_cxx_bridge.h header declares the C++ shim functions that wrap mlx::core operations. These shims handle the conversion from UniquePtr/reference semantics to the mlx::core::array value semantics used internally.

graph LR
    subgraph "Rust (mlxcel-core)"
        R[Model code<br/>e.g., Llama3Model] -->|ffi::matmul| B[cxx bridge<br/>lib.rs]
    end

    subgraph "C++ Shim Layer"
        B -->|mlx_cxx::matmul| SH[mlx_cxx_bridge.cpp]
    end

    subgraph "MLX C++ Core"
        SH -->|mlx::core::matmul| MLX[mlx::core]
        MLX -->|Metal GPU kernel| GPU[Apple GPU]
    end

Module Overview¶

Key Types¶

MlxArray¶

MlxArray is an opaque handle to an mlx::core::array. It is always owned via UniquePtr<MlxArray> in Rust, which corresponds to a heap-allocated wrapper holding the array value.

// Creating arrays
let zeros = ffi::zeros(&[batch, seq_len, hidden], ffi::FLOAT16);
let ones  = ffi::ones(&[heads, head_dim], ffi::FLOAT32);

// Operations return UniquePtr<MlxArray>
let result: UniquePtr<MlxArray> = ffi::matmul(&a, &b);

// Extract scalar
let value: f32 = ffi::item_f32(&logits);

// Get shape
let shape: Vec<i32> = ffi::array_shape(&arr);

MlxStream¶

MlxStream represents an MLX compute stream. Operations submitted to the same stream execute in order; operations on different streams may execute in parallel.

mlxcel uses a dedicated generation stream to overlap compute with memory transfers (lookahead pipelining):

let gen_stream = ffi::new_stream_on_device(true); // GPU stream
// Operations on gen_stream run asynchronously
let output = ffi::matmul_stream(&a, &b, &gen_stream);
ffi::synchronize_stream(&gen_stream); // wait for completion

KVCache¶

The KVCache struct holds the key and value matrices for one transformer layer:

pub struct KVCache {
    pub keys: Option<UniquePtr<MlxArray>>,   // [batch, heads, seq_len, head_dim]
    pub values: Option<UniquePtr<MlxArray>>,  // [batch, heads, seq_len, head_dim]
    pub offset: usize,                         // current sequence position
}

Core Modules¶

engine/ and generate.rs¶

generate.rs defines the LanguageModel trait every model implements. The decode loop lives in engine/: Engine owns the model and the CachePool that holds every open sequence's KV state, and every model forward, sampler draw, and finish step runs through it. The server's BatchScheduler drives the model only through Engine, and DirectEngine (engine/direct.rs) is the single-sequence raw-completion client that mlxcel generate, mlxcel-bench-decode, and the inference session use. A single sequence is a step batch of one; there is no separate single-sequence decode loop. Errors are per sequence, so one failing row finishes with an error while the rest of the batch keeps decoding.

flowchart TD
    A[Input tokens] --> B[Embed tokens]
    B --> C{First step?}
    C -->|Yes| D[Prefill: full prompt forward pass]
    C -->|No| E[Decode: single token forward pass]
    D --> F[Sample token from logits]
    E --> F
    F --> G{Stop condition?}
    G -->|EOS or max_tokens| H[Return generated text]
    G -->|No| I[Append token to input]
    I --> E

Key optimization: lookahead pipelining. When generating token n, the forward pass for token n+1 is submitted to the GPU stream before the CPU finishes processing the sampled token n. This hides the CPU-side sampling latency behind GPU execution. The engine exposes this as a split step (Engine::submit then Engine::finish_rows); DirectEngine and BatchScheduler both use it, and MLXCEL_FORCE_SYNC disables it.

layers.rs¶

Shared layer operations consumed by all model implementations:

• rms_norm(x, weight, eps): RMSNorm
• layer_norm(x, weight, bias, eps): LayerNorm
• linear(x, weight, bias): Linear projection
• rope(x, cos, sin, traditional): Rotary Position Embeddings
• attention(q, k, v, mask, scale): Scaled dot-product attention
• scaled_dot_product_attention(q, k, v, scale, mask): Flash-attention variant
• silu(x) / gelu(x) / relu(x): Activation functions
• take(x, indices, axis): Embedding lookup
• Various quantized matrix multiplication operations for 4-bit / 8-bit weights

sampling.rs¶

Implements the sampling stack in order:

1. Apply temperature scaling: logits / temperature
2. Apply logit bias (if provided)
3. Apply top-k filtering: keep only top-k logits
4. Apply top-p (nucleus) filtering
5. Apply min-p filtering
6. Apply XTC (Exclude Top Choices) sampling
7. Apply DRY (Don't Repeat Yourself) penalty
8. Apply repetition penalty
9. Apply frequency/presence penalties
10. Sample from the resulting distribution (or argmax for greedy)

weights.rs¶

Handles loading SafeTensors files:

// Load all .safetensors shards from a directory into a WeightMap
let weights = load_weights_from_dir(&model_path)?;

speculative/¶

Implements the drafters of speculative decoding. There is no draft-model decode loop of its own: the token-only drafters (PromptLookupDrafter, FullModelDrafter) are driven by the engine's Engine::verify_round, and the MTP and DFlash drafters have their own round loops (speculative/mtp/, drafter/). A round:

1. The drafter proposes up to k candidate tokens. FullModelDrafter runs a smaller checkpoint from its own KV caches and proposes its argmax; PromptLookupDrafter copies the continuation of an earlier occurrence of the sequence tail.
2. The target model verifies the current token and the proposals in a single forward pass.
3. A proposal is accepted when it equals the target's draw at that position (sampler-match; the argmax at temperature 0). See docs/speculative-acceptance.md for the rule each path runs.
4. The first rejected position emits the target's own draw, and the rejected proposals are trimmed from the target and drafter caches.
5. When most rounds accept more than one token, effective throughput rises.

Build System¶

The build system is in src/lib/mlxcel-core/build.rs. It builds MLX with CMake from src/lib/mlxcelverse/ and compiles the C++ shim layer against the resulting headers.

Key steps:

1. MLX build: runs the CMake project in src/lib/mlxcelverse/. CMakeLists.txt fetches upstream MLX at the pinned commit (the GIT_TAG of its FetchContent_Declare(mlx ...) block), copies the per-backend source overlays from overlays/{cpu,metal,cuda,rocm}/ over the fetched tree, and compiles libmlx.
2. Backend selection: on macOS, Metal and Accelerate are enabled by default (MLXCEL_BUILD_METAL=OFF turns Metal off); the cuda and rocm Cargo features select the CUDA and ROCm backends.
3. Compilation: compiles cpp/mlx_cxx_bridge.cpp and the other bridge sources, plus mlxcel's own kernel launchers from src/lib/mlxcelverse/kernels/, with -std=c++20.
4. Linking: links against the static libmlx.a and, on macOS, the Metal and Accelerate frameworks.

src/lib/mlxcel-core/
  build.rs              # Build script
  src/
    lib.rs              # cxx bridge declaration
    *.rs                # Rust modules
  cpp/
    mlx_cxx_bridge.h    # C++ header (declares shim functions)
    mlx_cxx_bridge.cpp  # C++ implementation (calls mlx::core)

To regenerate the cxx bindings after modifying lib.rs, run:

cargo build -p mlxcel-core

The cxx crate generates the necessary glue code at build time via a procedural macro; no separate code generation step is required.

Adding New Operations¶

To expose a new MLX operation to Rust:

1. Declare the shim function in cpp/mlx_cxx_bridge.h:

   namespace mlx_cxx {
       std::unique_ptr<MlxArray> my_new_op(
           const MlxArray& a,
           const MlxArray& b,
           float scale
       );
   }
   

2. Implement the shim in cpp/mlx_cxx_bridge.cpp:

   std::unique_ptr<MlxArray> mlx_cxx::my_new_op(
       const MlxArray& a,
       const MlxArray& b,
       float scale
   ) {
       auto result = mlx::core::my_new_op(a.arr, b.arr, scale);
       return std::make_unique<MlxArray>(std::move(result));
   }
   

3. Declare the binding in src/lib.rs:

   #[cxx::bridge(namespace = "mlx_cxx")]
   mod ffi {
       unsafe extern "C++" {
           // ...existing declarations...
           fn my_new_op(a: &MlxArray, b: &MlxArray, scale: f32) -> UniquePtr<MlxArray>;
       }
   }
   

4. Add a Rust wrapper in the appropriate module (e.g., src/lib/mlxcel-core/src/ops.rs):

   pub fn my_new_op(a: &MlxArray, b: &MlxArray, scale: f32) -> UniquePtr<MlxArray> {
       crate::ffi::my_new_op(a, b, scale)
   }
   

5. Add a test in src/lib/mlxcel-core/src/ffi_tests.rs:

   #[test]
   fn test_my_new_op() {
       let a = ffi::zeros(&[4, 4], ffi::FLOAT32);
       let b = ffi::ones(&[4, 4], ffi::FLOAT32);
       let result = crate::my_new_op(&a, &b, 1.0);
       assert_eq!(ffi::array_shape(&result), vec![4, 4]);
   }
   

When adding operations used by multiple model families, update the usage comment in src/lib/mlxcel-core/src/layers.rs to list all call sites. This comment serves as the re-test list for future changes to the operation.