[PROPOSAL] GPU-accelerated SeismicIndex build via cuSPARSE

[chishui]

Summary

Add optional GPU acceleration for the index-building (build()) path of
SeismicIndex, using NVIDIA cuSPARSE plus a small set of custom CUDA kernels.
Search stays entirely on CPU and is unchanged. On an NVIDIA L4 the GPU build is
~6.4× faster than a 32-core CPU (and ~15× faster than an 8-core CPU) on
the Big-ANN sparse SPLADE MS MARCO corpus, with cluster assignments and summary
vectors bit-identical to the CPU path on tie-free data.

GPU support is opt-in at compile time (-DNSPARSE_ENABLE_CUDA=ON, default
OFF) so CPU-only builds and CI are completely unaffected.

Motivation

Building a seismic index over large sparse corpora is expensive: on 8.84M
SPLADE documents it takes ~8 min on 32 cores and ~19 min on 8 cores. Index
building is an offline, embarrassingly-parallel, compute-heavy step — a natural
fit for GPU offload, while online search (latency-sensitive, small per-query
work) is best left on CPU. This mirrors FAISS's split, where GPU support is a
build/train accelerator rather than a search-path replacement for most index
types.

Scope

In scope: GPU offload of the two heaviest phases of build():

1. k-means document→centroid assignment (the dominant cost, ~87% of build).
2. summarize() per-term max-pool (the next-largest cost).

Out of scope (stays on CPU):

• search() — unchanged, always CPU.
• Inverted-list construction and the per-list prune (top-λ). These are a small
  fraction of build time (<1% for prune) and are branchy/sort-heavy — a poor
  GPU fit with low upside (see "Rejected alternatives").
• Non-float weights: the GPU path handles float (U32) only; U8/U16 quantized
  indices fall back to CPU automatically.

Background: where build() spends its time

SeismicIndex::build() → build_inverted_lists_clusters() runs, per inverted
list (there are ~30K), inside an OpenMP parallel loop:

1. prune_and_keep_doc_ids(λ) — keep the top-λ docs by weight (CPU sort).
2. RandomKMeans::train() → map_docs_to_clusters() — assign each doc to its
   nearest centroid. This is a sparse×dense matrix product.
3. summarize(α) — per cluster, per-term max-pool over the cluster's docs, then
   sort by value and truncate at the α mass boundary.

Measured phase split (thread-seconds across 32 threads, base_full, CPU build):

Phase	thread-seconds	share
prune	97.7	0.7%
assignment (train)	12,567.9	87.1%
summarize	1,764.0	12.2%

Assignment dominates, followed by summarize. Prune is negligible.

Design

1. A separate GPU module, gated at compile time

• New nsparse/gpu/gpu_cluster_assigner.{h,cu}. The header is always
  includable; the .cu compiles only when -DNSPARSE_ENABLE_CUDA=ON
  (defines NSPARSE_WITH_CUDA).
• New cmake/cuda.cmake discovers the CUDA toolkit + cuSPARSE and enables the
  CUDA language. Top-level option(NSPARSE_ENABLE_CUDA ... OFF).
• The CPU hook in map_docs_to_clusters() / summarize() is guarded by
  NSPARSE_WITH_CUDA and falls back to CPU automatically on any GPU error
  (OOM, no device, etc.).

This keeps the default build, existing CI, and non-CUDA platforms untouched.

2. Assignment as cuSPARSE SpMM (the core idea)

For one inverted list, stacking the candidate documents as a CSR matrix A
(n_docs × dim) and the cluster centroids as a dense matrix B
(dim × n_clusters) makes the assignment scores exactly C = A · B
(n_docs × n_clusters); each document's cluster is argmax over its row of
C. That is a textbook cusparseSpMM followed by a per-row argmax kernel.

The argmax kernel reproduces the CPU tie-break exactly (strict-greater update →
lowest cluster index wins), so results are bit-identical.

3. Keep the corpus resident; build inputs on-device

A naïve "one cuSPARSE call per list" offload is net-negative (it was 2.6×
slower than CPU in our first cut) for three reasons, all fixed in the design:

1. Resident corpus. The full corpus CSR is uploaded to the GPU once and
   cached (keyed by (pointer, n_vectors, nnz) — pointer alone is unsafe
   because a freed SparseVectors address can be reused). ~7 GB for base_full.
2. On-device gather/scatter. The per-list sparse A and dense B are built
   on the GPU from the resident corpus by small kernels, so only tiny doc-id /
   centroid-id arrays cross PCIe per list — not a 48 MB mostly-zero dense matrix.
3. Per-thread streams, no global lock. Each OpenMP worker thread gets its
   own cuSPARSE handle + CUDA stream + reusable (grow-on-demand) scratch
   buffers. This preserves the 32-way concurrency of the build loop; a global
   mutex would serialize all 32 threads through the one GPU.

4. summarize() max-pool: batch per list, not per cluster

summarize()'s per-term max-pool (67% of summarize) is offloaded; the
tie-order-sensitive sort/truncate stays on CPU so output ordering is preserved.

The critical design point is granularity. Clusters average only ~10 docs, so
a per-cluster GPU call (with its blocking sync) is dominated by launch/round-trip
overhead — a per-cluster version was ~5× slower than CPU (~27M blocking syncs
across the build). The shipped design issues one kernel launch per inverted
list — one thread block per cluster — collapsing ~27M syncs to ~30K. SPLADE
weights are ≥ 0, so a non-negative float compares monotonically as its
reinterpreted int32, enabling atomicMax on the accumulator. A per-(cluster, term) accumulator (n_clusters × dim) is reset only on touched slots so it is
reused across lists without a full re-zero.

5. Correctness model

• Assignment: bit-identical to the scalar CPU path (same math, same
  tie-break).
• summarize max-pool: bit-identical on tie-free data. The only possible
  difference vs. pure-CPU is at an exact float tie on the α-truncation
  boundary, where the (unstable) value-sort may keep a different term. This is
  rare and recall-neutral; the sort/truncate itself is unchanged CPU code.
• Parity is enforced by unit tests that assert GPU == CPU cluster membership and
  GPU == CPU max-pool (term, value, sum) (built only when CUDA is enabled).

6. Runtime controls

• NSPARSE_ENABLE_CUDA (CMake) — compile GPU support in. Default OFF.
• NSPARSE_GPU_MIN_DOCS — minimum list size to offload assignment (default
  1024). Small lists stay on CPU where they're cheaper.
• NSPARSE_GPU_SUMMARIZE — enable/disable summarize offload (default on when
  built with CUDA).
• NSPARSE_GPU_PROFILE / NSPARSE_BUILD_PROFILE / NSPARSE_SUMMARIZE_PROFILE
  — opt-in phase timers for diagnosis.

Benchmark

Method. benchmarks/index_build_benchmark.cpp times SeismicIndex::build()
only (data loaded once). GPU vs CPU selected at runtime via NSPARSE_GPU_MIN_DOCS
so both run from the same Release binary. Dataset: Big-ANN sparse vectors
(SPLADE MS MARCO) CSR. Params lambda=6000|beta=400|alpha=0.4.

Hardware. NVIDIA L4 (23 GB, sm_8.9); host with 32 vCPU / 121 GB. CUDA 13.

Build time: CPU vs GPU

Dataset	Docs / dim / nnz	CPU 8-core	CPU 32-core	GPU (L4)	Speedup vs 8-core	Speedup vs 32-core
base_small	100K / 30109 / 12.7M	82.0 s	30.6 s	6.6 s	12.4×	4.6×
base_full	8.84M / 30109 / 1.12B	1163.8 s	483 s	74.9 s	15.5×	6.4×

Speedup improves with scale (fixed overhead amortized, larger per-list
matmuls). All builds verified bit-identical to CPU; full test suite (365 tests)
passes with CUDA enabled.

Phase breakdown (base_full, thread-seconds)

Phase	CPU build	GPU build
prune	97.7	88.7
assignment	12,567.9	148.8 (≈84× less)
summarize	1,764.0	759.7 (batched GPU max-pool)

After offload, the single L4 becomes the shared bottleneck between assignment
and summarize (threads ~37% blocked on the GPU) — the remaining ceiling on one
GPU.

Context vs dense GPU build (for calibration)

The dense k-NN CAGRA build (opensearch-project/k-NN#2293) reports 30–50×
vs an 8-vCPU baseline. Sparse Seismic build here is 12–15× vs 8-core. The
gap is structural, not an optimization deficit: sparse assignment has far lower
arithmetic intensity (scattered gathers, ~127 nnz over a 30K-dim space vs dense
contiguous vectors), and some build phases (prune, summarize sort) are branchy
and stay on CPU (Amdahl's law).

Rejected / deferred alternatives

• Custom fused SpMM+argmax kernel (replace cuSPARSE): implemented and tested;
  slightly slower than cuSPARSE on base_full's 4400×440 problems (cuSPARSE's
  tiled algorithm has better B-matrix cache behavior). Kept cuSPARSE.
• GPU pruning / full-GPU build: prune is <1% of build time and is a
  sort+top-k per list — poor GPU fit, negligible upside. Deferred.
• Per-cluster summarize offload: ~5× slower due to per-cluster sync storm;
  replaced by the per-list batched kernel above.

Rollout / risks

• Opt-in only. Default build unchanged; no new required dependency.
• Build reproducibility: must be compiled Release (-O3). An
  unoptimized build is ~2–3× slower on both CPU and GPU paths.
• GPU memory: the resident corpus (~7 GB for base_full) plus per-thread
  summarize accumulators (2 · n_clusters · dim · 4 B per thread) must fit in
  device memory; the summarize accumulator footprint scales with worker-thread
  count and should be capped on smaller GPUs (fall back to CPU on OOM — already
  handled).
• Numerics: bit-identical on tie-free data; documented exact-tie edge case
  at the summarize truncation boundary.

Proposed implementation (already prototyped)

Reference implementation on branch gpu-cusparse-build-acceleration:

• d8b0a1d — GPU assignment via cuSPARSE (resident corpus, on-device
  gather/scatter, per-thread streams) + CMake wiring + parity tests.
• e871166 — opt-in profiling + env-configurable benchmark iterations.
• 81c51f5 — batched-per-list GPU summarize max-pool.

Happy to open a PR against these once the design is agreed.