ExMaxsimCpu

High-performance MaxSim (Maximum Similarity) scoring for Elixir using BLAS GEMM operations with SIMD acceleration. This is a port of the maxsim-cpu package by Mixedbread.ai.

What is MaxSim?

MaxSim is a scoring function used in ColBERT-style neural retrieval models. For each query token, it finds the maximum similarity with any document token, then sums these maximums to produce the final score:

MaxSim(Q, D) = Σᵢ maxⱼ(Qᵢ · Dⱼ)

This library provides a highly optimized implementation using:

• BLAS GEMM for efficient matrix multiplication (OpenBLAS on Linux, Accelerate on macOS)
• SIMD for fast max reduction (AVX2 on x86_64, NEON on ARM64)
• Rayon for parallel document processing
• Optional libxsmm support for additional performance on Intel CPUs

Installation

Add ex_maxsim_cpu to your dependencies in mix.exs:

def deps do
  [
    {:ex_maxsim_cpu, "~> 0.1.0"},
    {:nx, "~> 0.7"}  # Required for tensor operations
  ]
end

System Requirements

• macOS (Apple Silicon): Uses Accelerate framework (no additional setup needed)
• Linux x86_64: Requires OpenBLAS; AVX2 is recommended for SIMD speedups

  # Ubuntu/Debian
  sudo apt-get install libopenblas-dev
  
  # Fedora/RHEL
  sudo dnf install openblas-devel

  For source builds that should use AVX2, set:

  export RUSTFLAGS="-C target-feature=+avx2"

Rust Toolchain

This library uses Rustler to compile native code. Ensure you have Rust installed:

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

Usage

With Nx Tensors (Recommended)

# Create query tensor: [q_len, dim]
query = Nx.tensor([
  [0.1, 0.2, 0.3, 0.4],
  [0.5, 0.6, 0.7, 0.8]
], type: :f32)

# Create documents tensor: [n_docs, d_len, dim]
docs = Nx.tensor([
  [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]],
  [[0.9, 0.8, 0.7, 0.6], [0.5, 0.4, 0.3, 0.2]]
], type: :f32)

# Compute MaxSim scores
scores = ExMaxsimCpu.maxsim_scores(query, docs)
# => #Nx.Tensor<f32[2]>

Note: MaxSim expects L2-normalized embeddings. Normalize query and docs along the token axis before scoring.

Variable-Length Documents

query = Nx.tensor([[1.0, 0.0], [0.0, 1.0]], type: :f32)

# Documents with different lengths
doc1 = Nx.tensor([[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], type: :f32)  # 3 tokens
doc2 = Nx.tensor([[0.5, 0.5]], type: :f32)                          # 1 token

scores = ExMaxsimCpu.maxsim_scores_variable(query, [doc1, doc2])

Performance Tuning

Environment Variables

• RAYON_NUM_THREADS: Control Rayon thread pool size (default: number of CPUs)
• OPENBLAS_NUM_THREADS=1: Recommended to avoid oversubscription with Rayon

Example Configuration

export OPENBLAS_NUM_THREADS=1
export RAYON_NUM_THREADS=8

libxsmm Support (Optional)

For additional performance on Intel CPUs, you can enable libxsmm:

1. Build libxsmm:

   git clone https://github.com/libxsmm/libxsmm.git
   cd libxsmm
   make -j$(nproc)
   export LIBXSMM_DIR=$(pwd)

2. Build with the feature enabled:

   cd native/maxsim_cpu
   cargo build --release --features use-libxsmm

Benchmarks

ExMaxsimCpu delivers strong CPU performance, especially on small-to-moderate batch sizes:

Benchmark Environment

• OS: macOS 26.2 (25C56)
• CPU: Apple M4 Pro (12 physical / 12 logical cores)
• Architecture: arm64
• Elixir: 1.19.5 (OTP 28)
• BLAS: Accelerate
• Env: OPENBLAS_NUM_THREADS=1, RAYON_NUM_THREADS=8
• Nx CPU backend: Torchx (CPU)

Performance Summary

Notes:

• Nx BinaryBackend is an unaccelerated baseline; EXLA/Torchx CPU will be faster.
• Nx CPU backend results (EXLA or Torchx) are included when available.
• MPS timings are for small shapes where transfer overhead can dominate.
• Results are machine- and shape-dependent; run your own benchmarks for production sizing.

Performance Comparison

Benchmark on Apple Silicon (this run), showing ExMaxsimCpu (blue), Nx BinaryBackend (orange), and optional Nx CPU/MPS backends on a log scale.

Speedup Analysis

Why can CPU win on these shapes?

1. Optimized BLAS: Uses Apple Accelerate (or OpenBLAS on Linux) which is highly tuned for the CPU architecture
2. SIMD Instructions: Hand-optimized AVX2/NEON code for max-reduction operations
3. No Transfer Overhead: GPU implementations incur memory transfer costs
4. Cache Efficiency: Tiled algorithms keep data in CPU cache

Detailed Comparisons

Running Benchmarks

Generate benchmark data (includes optional MPS and Nx CPU backends when available):

OPENBLAS_NUM_THREADS=1 mix run bench/generate_plots.exs

Generate plots:

uv run bench/plot_benchmarks.py

Note: For MPS benchmarks, ensure torchx is installed. For Nx CPU backends, install torchx or add exla to your dev deps. The benchmark auto-detects availability.

Platform Support

Credits

This is the Elixir binding for maxsim-cpu, a high-performance MaxSim implementation in Rust.