CUDA & C++ 20 Compatibility, DLPack, Float8 & 3D Mesh Speedups

[ashvardanian]

CUDA & C++ 20 Compatibility

NVCC 13 caps its language-standard flag at C++20, and our multi-argument subscript overloads from C++23 P2128 made tensor.hpp unparseable by cudafe++. We added call-operator primaries that mirror every multi-argument subscript overload in the tensor view, span, and owning container types, and kept the bracket sugar behind an __cpp_multidimensional_subscript feature test so older toolchains pick the portable spelling automatically. Downstream CUDA callers now parse the tensor header without touching their language-standard flag.

DLPack & Zero-copy Exchange with PyTorch, JAX, & Arrow

Tensors now exchange zero-copy in both directions with every Python framework that implements the DLPack protocol — PyTorch, NumPy, JAX, CuPy, TensorFlow, PyArrow, MLX, ONNX Runtime, TVM, MXNet, NNabla — using DLPack 1.3's versioned capsules and the max_version handshake. This finally carries semantic dtype identity across the bridge: bf16 and the four narrow float variants E4M3FN, E5M2, E2M3, and E3M2 round-trip without losing their type, where PEP 3118 and the legacy array interface previously degraded them to raw unsigned bytes. The importer accepts every device whose pointer is host-dereferenceable — plain CPU, pinned host memory on CUDA and ROCm, CUDA managed unified memory, Intel oneAPI host and shared USM, and Metal on Apple Silicon — while pure device memory is rejected with the offending device code named. The exporter stays strict and only emits the CPU device. Sub-byte types u1, u4, and i4 ride as byte containers, and the ABI is declared inline as six structs and twelve device codes rather than vendoring an external header, mirroring NumPy's own approach. Validated against torch 2.11, numpy 2.4, jax 0.10, tensorflow 2.21, pyarrow 23, cupy 13.6, and onnxruntime 1.24 on H100 with 127 tests passing.

import numpy as np, torch, numkong as nk

# NumKong → PyTorch: zero-copy FP8 round-trip preserves dtype identity.
src = torch.zeros(4, 6, dtype=torch.float8_e4m3fn)
nk_view = nk.from_dlpack(src); pt_back = torch.from_dlpack(nk_view)
assert nk_view.shape == (4, 6) and nk_view.dtype == "e4m3" and pt_back.dtype == torch.float8_e4m3fn

# Mutation through one view is visible through the other — proves zero-copy.
tensor = nk.Tensor(np.arange(24, dtype=np.float32).reshape(4, 6)); pt = torch.from_dlpack(tensor)
pt[0, 0] = 99; assert np.asarray(tensor)[0, 0] == 99

Upstream DLPack PRs that this bridge interoperates with, already referenced from our DLPack interop source:

• pytorch/pytorch#57110 — PyTorch DLPack protocol, later upgraded to DLPack 1.0 with the max_version handshake
• tensorflow/community#180 — TensorFlow DLPack RFC, implemented as tf.experimental.dlpack
• microsoft/onnxruntime#23110 — ONNX Runtime OrtValue DLPack enabled by default for inference

Faster Single-Pass 3D Mesh Alignment Algorithms

Kabsch/Umeyama mesh alignment now folds into a single pass via the trace identity

$$\mathrm{SSD} = \lVert a - \bar a \rVert^2 + \lVert b - \bar b \rVert^2 - 2,\mathrm{tr}(R \cdot H)$$

replacing the earlier two-pass approach — covariance first, transformed-SSD second — across all nine backends. An identity-dominant short-circuit skips SVD entirely when $H$ approximates a positive diagonal, saving around 500 cycles on already-aligned inputs. Two new backends land alongside: a Genoa kernel that uses VDPBF16PS for channel-grouped bf16 reductions, and a NEON+FP16FML kernel that uses vfmlalq for fp16 widening FMA, while the existing NEON+BFDOT path picks up vbfdotq_f32 for its bf16 stats pass. A centered-RMSD bug in the NEON and NEON+FP16FML paths is fixed in passing.

Faster Float8 Linear Algebra on x86

Pairwise FP8 distance kernels — sqeuclidean, euclidean, and angular — on Skylake and Haswell now compute the squared difference directly in F32 after a free-shift widen. E5M2 abuses its shared exponent bias with F16: a byte-to-word unpack against zero places the byte as a valid F16 encoding. E4M3 uses a Giesen-style fake-F16 cast that shifts the mantissa up by seven, reinjects the sign at bit fifteen, widens with vcvtph2ps, and multiplies by 256 to correct the bias delta. Per-pair speedups range from 1.4× for E4M3 angular on Skylake to 4.9× for E5M2 sqeuclidean on Haswell on a pinned Xeon 6776P. The redundant Genoa E5M2 pairwise kernels are deleted because the rewritten Skylake path runs on Genoa silicon and beats the old vdpbf16ps-chain form by 2.4×.

Stateful FP8 GEMMs follow the same trajectory. E5M2 byte-packs into a new dtype-specific update helper that runs two FMA chains into a single state accumulator, landing at 1.4–2.5× on Skylake and up to 3.2× on Haswell for the packed dot, angular, and euclidean variants. E4M3 GEMMs on Skylake switch to an asymmetric F16-pack scheme where A streams as F32 while B is pre-cast at pack time and stored as F16, halving packed-B memory with compute neutral against baseline. Granite Rapids gets a brand-new E5M2 GEMM that packs E5M2 into F16 with a single byte shift and runs TDPFP16PS over F16 tiles, beating the Sapphire AMX BF16 path on E5M2 inputs with better intermediate precision at the same throughput. Dispatch wires it ahead of Sapphire AMX so Granite hardware automatically picks it up.