Technology · Discipline

Approximate nearest-neighbour search

I built exact, inverted-file and lattice nearest-neighbour search as GPU kernels with CPU references and equivalence tests, inside a vector store whose index builds are deterministic and whose search reports its own health.

CUDA from Rust, raw Driver APINative SM 120 binariesNVIDIA Blackwell
Three search families · one top-k engine
Shared top-k engine
  • 0.85IVF recall at 10 against exact search, the floor a CI test holds
  • 128largest k; cosine, dot product and L2 are entry points in one binary
  • 1e-4relative tolerance between each GPU executor and its CPU reference

I write first-party GPU search when the available libraries cannot meet the platform, and I hold it to exact references and CI recall floors from the first kernel.

Why I wrote the search myself

Inference Fabric produced bulk embedding and hidden-state vectors that I wanted searched and analysed on the GPU where they were stored. The established GPU vector libraries gave me search and whole-dataset clustering, and could not reach the platform. I wrote the search myself: a CUDA engine driven from Rust through the raw CUDA Driver API, compiled for the NVIDIA Blackwell GPU generation.

CUDA 13 on Blackwell, from Rustbytes handed to my kernelsEstablished GPU vector librariesRust CUDA wrappers of the timeMy engine on the raw Driver API
A blank cell is a question the option never reached. No PTX fallback.

The exact kernel

The exact kernel scores every vector against the query and never materialises a score matrix.

The database, streamed in tiles

A 16-row slice moves through shared memory with the query, 64 dimensions at a time.

One block of 256 threads · one query
Shared buffer · best k
only the top k per query is written out

One top-k engine for every family

All search families share one top-k engine. It accepts either a score row per query or a set of partial top-k lists, and merges in two passes through a workspace. I kept cooperative launches and grid-wide synchronisation out of it. Ties break on the candidate's position, so the same inputs always produce the same ranking.

Pass 1 · sorted partial top-k per block
Pass 2 · one block per query

ties broken on position

A fix to selection lands in every family at once.

IVF-flat is the first approximate family. At query time one block per query walks every centroid, keeps the best nprobe cells, and scores their posting lists through the shared top-k engine. Seeding on the CPU makes its index builds reproducible bit for bit.

illustrative points; two cells shaded as the nprobe cells for one query

build
on the GPU: distances, then assign and update
on the host, from the build seed: the weighted random draw
The configuration, training set and result are hashed, and an index can be rebuilt from its inputs to check it; an empty cell keeps its previous centroid.

Lattice search runs the same idea in two levels over the clustered structures the distillation pipelines build. The first level picks the nearest cluster centroids for each query. The second launches one block per query and selected cluster, scores that cluster's members, and writes a partial top-k in the exact layout the shared engine merges. There is no lattice-specific reducer to maintain.

Held to references in CI

Every executor has a CPU reference implementation and equivalence tests, with a relative tolerance wide enough that f32 reduction order does not cause false failures. A CI integration test holds IVF recall against exact search on clustered data, so a recall regression fails the build.

Recall against exact search
the mark is the floor a CI test holds
CPU reference
GPU executor
equal within tolerance
What CI checks
  • GPU result against its CPU reference, within toleranceequivalent
  • race-freedom and cross-block correctnesstested
  • recall rises with nprobetested
  • IVF recall below the floorbuild fails

Exact search is the reference that every approximate family is measured against in CI.

Inside the vector store

The Vector Shard Manager is the store the search runs inside, and storage and kernels change independently. It reports the health of the clusters a query searched alongside the results. This engine replaced an earlier path that ran FAISS from Python over vectors pulled through SQL.

Shard Manager daemonlinks no CUDA
Type-only interface cratekernels see only leased batch types
GPU executors
exact, IVF, lattice
A CI check fails the build if the daemon ever links CUDA.