NumKong: Mixed Precision for All

Portable mixed-precision math, linear-algebra, & retrieval library with 2'000+ SIMD kernels for x86, Arm, RISC-V, LoongArch, Power, & WebAssembly, leveraging rare algebraic transforms with both 1D & 2D registers like AMX & SME, covering 15+ numeric types from 4-bit integers & 6-bit floats to 128-bit complex numbers, validated against 118-bit extended-precision baselines with saturation, casting, & rounding edge-case coverage, in a 5-100x smaller binary than other BLAS-like alternatives, co-designed with Tensor abstractions in C++, Python, Rust, JavaScript, GoLang, & Swift.

Latency, Throughput, & Numerical Stability

Most libraries return dot products in the same type as the input — Float16 × Float16 → Float16, Int8 × Int8 → Int8. This leads to quiet overflow: a 2048-dimensional i8 dot product can reach ±10 million, but i8 maxes out at 127. NumKong promotes to wider accumulators — Float16 → Float32, BFloat16 → Float32, Int8 → Int32, Float32 → Float64 — so results stay in range.

Single 2048-d dot product on Intel Sapphire Rapids, single-threaded. Each cell shows gso/s, mean relative error vs higher-precision reference. gso/s = Giga Scalar Operations per Second — a more suitable name than GFLOP/s when counting both integer and floating-point work. NumPy 2.4, PyTorch 2.10, JAX 0.9.

A fair objection: PyTorch and JAX are designed for throughput, not single-call latency. They lower execution graphs through XLA or vendored BLAS libraries like Intel MKL and Nvidia cuBLAS. So here's the same comparison on a throughput-oriented workload — matrix multiplication:

Matrix multiplication (2048 × 2048) × (2048 × 2048) on Intel Sapphire Rapids, single-threaded. gso/s = Giga Scalar Operations per Second, same format. NumPy 2.4, PyTorch 2.10, JAX 0.9, same versions.

For f64, compensated "Dot2" summation reduces error by 10–50× compared to naive Float64 accumulation, depending on vector length. For f32, widening to Float64 gives 5–10× lower error. The library ships as a relatively small binary:

Every kernel is validated against 118-bit extended-precision baselines with per-type ULP budgets across log-normal, uniform, and Cauchy input distributions. Tests check triangle inequality, Cauchy-Schwarz bounds, NaN propagation, overflow detection, and probability-simplex constraints for each ISA variant. Results are cross-validated against OpenBLAS, Intel MKL, and Apple Accelerate. A broader throughput comparison is maintained in NumWars.

Quick Start

What's Inside

NumKong covers 17 numeric types — from 6-bit floats to 64-bit complex numbers — across dozens of operations and 30+ SIMD backends, with hardware-aware defaults: Arm prioritizes f16, x86 prioritizes bf16.

Language Bindings

Design Decisions

• Avoid loop unrolling and scalar tails.
• Don't manage threads and be compatible with any parallelism models.
• Don't manage memory and be compatible with arbitrary allocators & alignment.
• Don't constrain ourselves to traditional BLAS-like Matrix Multiplication APIs.
• Don't throw exceptions and pass values by pointers.
• Prefer saturated arithmetic and avoid overflows, where needed.
• Cover most modern CPUs with flexible dispatch and wait for them to converge with GPUs.

The rest of this document unpacks the functionality and the logic behind the design decisions.

Auto-Vectorization & Loop Unrolling

Most "optimized SIMD code" is a 2–4x unrolled data-parallel for-loop over f32 arrays with a serial scalar tail for the last few elements:

float boring_dot_product_f32(float const *a, float const *b, size_t n) {
    __m256 sum0 = _mm256_setzero_ps(), sum1 = _mm256_setzero_ps();
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        sum0 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i), sum0);
        sum1 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i + 8), _mm256_loadu_ps(b + i + 8), sum1);
    }
    float result = _mm256_reduce_add_ps(_mm256_add_ps(sum0, sum1));
    for (; i < n; i++) result += a[i] * b[i]; // serial tail
    return result;
}

This kind of unrolling has been a common request for NumKong, but the library avoids it by design.

Modern CPUs already "unroll" in hardware. Out-of-order engines with reorder buffers of 320–630 entries (Zen 4: 320, Golden Cove: 512, Apple Firestorm: ~630) can keep a dozen of loop iterations in-flight simultaneously. The physical register file is much larger than the ISA-visible architectural registers — Skylake has ~180 physical integer registers behind 16 architectural GPRs, and ~168 physical vector registers behind 32 architectural ZMMs. The register renaming unit maps the same zmm0 in iteration N and iteration N+1 to different physical registers, extracting cross-iteration parallelism automatically — exactly the benefit that source-level unrolling was historically supposed to provide.

Unrolling works against NumKong's goals. Every unrolled copy is a distinct instruction in the binary. With 1,500+ kernel endpoints across 30+ backends, even 2x unrolling would inflate the .text section by megabytes — directly impacting install size for Python wheels, NPM packages, and Rust crates. Larger loop bodies also increase instruction-cache and micro-op-cache pressure; Agner Fog also recommends:

"avoid loop unrolling where possible in order to economize the use of the micro-op cache".

A loop that spills out of the uop cache falls back to the slower legacy decoder, making the "optimized" version slower than the compact original. For a header-only library, unrolling also compounds compilation time: register allocation is NP-hard (reducible to graph coloring), and unrolling multiplies the number of simultaneously live ranges the allocator must consider, increasing compile time super-linearly across every translation unit that includes the headers.

Serial tails are a correctness hazard. The leftover elements after the last full SIMD chunk run through a scalar loop that silently drops FMA fusion, compensated accumulation, and saturating arithmetic — producing results with different numerical properties than the SIMD body. NumKong often uses masked loads instead (_mm512_maskz_loadu_ps on AVX-512, predicated svld1_f32 on SVE), processing every element through the same arithmetic path regardless of alignment. It's not exactly orthogonal to loop-unrolling, but makes a different kernel layout more compatible.

The gains come from elsewhere. On Intel Sapphire Rapids, NumKong was benchmarked against auto-vectorized code compiled with GCC 12. GCC handles single-precision float well, but struggles with _Float16 and other mixed-precision paths:

NumKong's f16 kernels are faster than GCC's f32 output — not because of unrolling, but because they use F16C conversion instructions, widening FMA pipelines, and compensated accumulation that compilers do not synthesize from a plain for loop. The same story repeats for bf16, e4m3, i8, and i4: these types require algorithmic transformations — lookup tables, algebraic domain shifts, asymmetric VNNI tricks — that live beyond the reach of auto-vectorization.

Parallelism & Multi-Threading

BLAS libraries traditionally manage their own thread pools. OpenBLAS spawns threads controlled by OPENBLAS_NUM_THREADS, Intel MKL forks its own OpenMP runtime via MKL_NUM_THREADS, and Apple Accelerate delegates to GCD (Grand Central Dispatch). This works in isolation — but the moment your application adds its own parallelism (joblib, std::thread, Tokio, GCD, OpenMP), you get thread oversubscription: MKL spawns 8 threads inside each of your 8 joblib workers, producing 64 threads on 8 cores, thrashing caches and stalling on context switches. The Python ecosystem has built entire libraries just to work around this problem, and scikit-learn's documentation devotes a full page to managing the interaction between joblib parallelism and BLAS thread pools.

NumKong takes a different position: the numerics layer should not own threads. Modern hardware makes the "spawn N threads and split evenly" model increasingly untenable:

• Server-grade CPUs have hundreds of cores split across sockets, chiplets, and tiles, resulting in dozens of physical NUMA domains with vastly different memory access latencies. A thread pool that ignores NUMA topology will spend more time on remote memory stalls than on arithmetic.
• Consumer-grade CPUs pack heterogeneous Quality-of-Service core types on the same die — Intel P-cores and E-cores run at different frequencies and sometimes support different ISA extensions. A naive work-split gives equal chunks to fast and slow cores, and the whole task stalls waiting for the slowest partition.
• Real-time operating systems in robotics and edge AI cannot afford to yield the main thread to a BLAS-managed pool. These systems need deterministic latency, not maximum throughput.

Instead, NumKong exposes row-range parameters that let the caller partition work across any threading model. For GEMM-shaped dots_packed, this is straightforward — pass a slice of A's rows and the full packed B to compute the corresponding slice of C. For SYRK-shaped dots_symmetric, explicit start_row / end_row parameters control which rows of the symmetric output matrix a given thread computes. The GIL (Global Interpreter Lock) is released around every kernel call, making NumKong compatible with concurrent.futures, multiprocessing, or any other parallelism model:

import concurrent.futures, numkong as nk, numpy as np
vectors, num_threads = np.random.randn(1000, 768).astype(np.float32), 4
output = nk.zeros((1000, 1000), dtype="float32")
def compute_slice(t):
    start = t * (len(vectors) // num_threads)
    end = start + len(vectors) // num_threads if t < num_threads - 1 else len(vectors)
    nk.dots_symmetric(vectors, out=output, start_row=start, end_row=end)
with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as pool:
    list(pool.map(compute_slice, range(num_threads)))

For users who want a ready-made low-latency thread pool without the oversubscription baggage of OpenMP, we built ForkUnion — a minimalist fork-join library for C, C++, and Rust that avoids mutexes, CAS atomics, and dynamic allocations on the critical path, with optional NUMA pinning on Linux.

Memory Allocation & Management

BLAS libraries typically allocate internal buffers during GEMM — OpenBLAS packs matrices into L2/L3-sized panels via per-thread buffer pools backed by mmap or shmget. This hidden allocation has caused real problems: 14 lock/unlock pairs per small GEMM call throttling 12-thread scaling to 2x, silently incorrect results from thread-unsafe allocation in np.dot, and deadlocks after fork() due to mutex state not being reset in child processes. The BLASFEO library was created specifically for embedded model-predictive control where malloc during computation is unacceptable.

NumKong never allocates memory. Following the same philosophy as Intel MKL's packed GEMM API (cblas_sgemm_pack_get_size → cblas_sgemm_pack → cblas_sgemm_compute), NumKong exposes typed three-phase interfaces — nk_dots_packed_size_* → nk_dots_pack_* → nk_dots_packed_* — where the caller owns the buffer and NumKong only fills it.

The reason GEMM libraries repack matrices at all is that every hardware target has a different preferred layout — Intel AMX expects B in a VNNI-interleaved tile format (pairs of BFloat16 values packed into DWORDs across the K dimension), while Arm SME wants column vectors for its FMOPA outer-product instructions. Since GEMM is