startrail-org--leann/docs/flashlib_backend_guide.md

FlashLib Backend Guide

LEANN now ships an optional, GPU-accelerated search backend powered by FlashLib — a library of classical machine-learning operators built on Triton and CuteDSL. This guide covers what it is, when to use it, how to install it, and how it plugs into LEANN.

What is FlashLib?

FlashLib (pip install flashlib) is a GPU library of classical ML primitives (k-means, DBSCAN, PCA, SVD, UMAP, t-SNE, IVF-Flat ANN, and more). LEANN uses its IVFFlat index — an inverted-file flat approximate-nearest-neighbor (ANN) index that runs entirely on CUDA tensors. At a fixed (nlist, nprobe) it probes the same candidate set as a reference IVF-Flat (FAISS / cuVS), so recall is predictable, while the search itself is accelerated on the GPU.

import torch
from flashlib import IVFFlat

db = torch.randn(1_000_000, 128, device="cuda")
index = IVFFlat(nlist=1024, nprobe=16).fit(db)
distances, indices = index.kneighbors(torch.randn(10_000, 128, device="cuda"), n_neighbors=10)

When to use it

Backend	Best for	Storage	Hardware
hnsw (default)	Laptop / CPU, max storage savings via recomputation	~3% of raw (pruned graph)	CPU
diskann	Larger-than-memory datasets	On-disk graph	CPU
ivf	Incremental add/remove without rebuild	Full vectors (FAISS)	CPU
flashlib	High-throughput search on a CUDA GPU	Full vectors (.npy)	CUDA GPU
flashlib_ivf	GPU IVF-Flat (approximate) — the GPU counterpart of ivf	Full vectors (.pt)	CUDA GPU

Use FlashLib when you already have a GPU and want fast IVF-Flat search. It stores the full float32 vectors rather than a pruned graph, so it trades LEANN's storage savings for raw GPU search speed.

Requirements

• A CUDA GPU (required at search time; building the index only needs numpy).
• flashlib and torch (installed automatically with the extra below).

Installation

The backend is optional — it is not pulled in by a default LEANN install.

# From a LEANN source checkout
uv sync --extra flashlib

# Or as a standalone package
pip install leann-backend-flashlib

LEANN auto-discovers any installed leann-backend-* package, so once it is installed the flashlib backend name is available with no further configuration.

Usage

Python API

from leann import LeannBuilder, LeannSearcher

builder = LeannBuilder(backend_name="flashlib")   # nlist=1024, distance_metric="mips"
builder.add_text("LEANN recomputes embeddings on the fly to cut storage by ~97%.")
builder.add_text("FlashLib runs IVF-Flat search on the GPU.")
builder.build_index("demo.leann")

searcher = LeannSearcher("demo.leann")
results = searcher.search("How does LEANN save storage?", top_k=3)
for r in results:
    print(r.score, r.text)

Example apps / CLI

source .venv/bin/activate
python -m apps.document_rag \
    --query "What are the main techniques LEANN explores?" \
    --backend-name flashlib

How it works

FlashLib's IVFFlat builds its index in GPU memory and has no on-disk format. The LEANN backend bridges that gap:

1. Build (FlashlibBuilder): persists the raw float32 vectors as <index>.flashlib.npy and an id map as <index>.flashlib_id_map.json.
2. Search (FlashlibSearcher): loads the vectors into a CUDA tensor and reconstructs the index once via IVFFlat(nlist, nprobe).fit(db) at start-up, then answers every query with index.kneighbors(...).

FlashLib's only distance metric is squared L2. For mips / cosine the backend L2-normalizes both the database and the query vectors; on unit vectors, squared-L2 ranking is equivalent to inner-product / cosine ranking, so results match the other backends. nlist is clamped to the corpus size (the k-means constraint), and nprobe is derived from the search complexity knob.

Parameters

Parameter	Default	Meaning
nlist (build)	1024	Number of IVF partitions; clamped to the number of vectors.
distance_metric (build)	"mips"	mips, cosine, or l2.
nprobe (search)	derived from complexity	Partitions probed per query — the recall knob (higher = better recall, slower).

FlashLib IVF backend (flashlib_ivf)

The flashlib_ivf backend (leann-backend-flashlib-ivf) is the GPU counterpart of LEANN's FAISS ivf backend: an IVF-Flat (inverted file) index that coarse-quantizes the corpus into nlist cells with k-means and, at search time, scans only the nprobe nearest cells — entirely on CUDA tensors via FlashLib's flash_ivf_flat. At a fixed (nlist, nprobe) the GPU and CPU IVF probe nearly the same candidate set, so recall is comparable; the only difference is GPU vs CPU kernels. nprobe is the recall knob (defaults to min(complexity, nlist)).

uv sync --extra flashlib-ivf          # or: pip install leann-backend-flashlib-ivf

from leann import LeannBuilder, LeannSearcher

builder = LeannBuilder(backend_name="flashlib_ivf", nlist=4096, distance_metric="cosine")
builder.add_text("LEANN recomputes embeddings on the fly to cut storage by ~97%.")
builder.build_index("demo.leann")

searcher = LeannSearcher("demo.leann")
results = searcher.search("How does LEANN save storage?", top_k=10, complexity=32)  # nprobe=32

How it works: the builder trains the coarse quantizer on the GPU and persists the built index tensors with torch.save (<index>.flashlib_ivf.pt) plus an id map (<index>.flashlib_ivf_id_map.json); the searcher reloads them onto the GPU once (no k-means re-train). A CUDA GPU is required at both build (k-means) and search time. FlashLib IVF ranks by squared L2, so mips/cosine L2-normalize the vectors (squared-L2 ranking then matches inner-product/cosine).

Speed — IVF (GPU) vs IVF (CPU)

python benchmarks/flashlib_ivf_vs_faiss_ivf.py \
    --sizes 100000 1000000 --nprobe-sweep 1 8 32 128 --cpu-threads 8

flashlib_ivf (GPU) vs the FAISS ivf backend (CPU) at the same nlist, sweeping nprobe (NVIDIA H200, faiss-cpu at 8 threads, 768-dim, top-k=10). GPU latency stays ~flat while CPU latency grows linearly with nprobe, so the GPU lead widens exactly as you raise recall:

Corpus	nprobe	GPU lat	CPU lat	GPU q/s	CPU q/s	Recall (GPU/CPU)	Speedup (lat / tpt)
1M	8	0.45 ms	1.14 ms	107k	5.9k	0.340 / 0.321	2.6× / 18×
1M	32	0.46 ms	3.00 ms	141k	1.9k	0.400 / 0.350	6.5× / 75×
1M	128	0.55 ms	9.91 ms	95k	0.6k	0.539 / 0.423	18× / 159×

Build (1M, nlist=4096): GPU 10.6 s vs FAISS CPU 140.7 s — a 13× faster build (GPU k-means vs CPU k-means training).

Honest caveats: at very low nprobe (e.g. 1) single-query GPU latency (~0.44 ms) is higher than CPU, because the per-query work is tiny and GPU kernel-launch overhead dominates; the GPU advantage grows with nprobe (higher recall), batch size, and corpus size. The absolute recall above is low because the synthetic mixture-of-Gaussians corpus has more clusters than nlist; on real embeddings recall is far higher — this benchmark isolates the GPU-vs-CPU relative comparison at matched (nlist, nprobe).

Notes & limitations

• GPU-only at search time; building an index works on a CPU-only machine.
• Stores full vectors, so it does not benefit from LEANN's graph-pruning storage savings — pick hnsw if minimizing disk footprint is the priority.
• Query embeddings are computed through the standard LEANN embedding path (the HNSW ZMQ embedding server when available, otherwise direct model loading), so any LEANN-supported embedding model works.