Gpu cusparse build acceleration#18
Open
chishui wants to merge 8 commits into
Open
Conversation
Accelerate the index-building (build()) path with NVIDIA cuSPARSE while
keeping search() entirely on CPU. The k-means document->centroid
assignment inside build() is a sparse(docs CSR) x dense(centroids)
matrix product, which maps onto cusparseSpMM plus a per-row argmax.
Implementation (nsparse/gpu/gpu_cluster_assigner.{h,cu}):
- Full corpus CSR uploaded to the GPU once, cached by (pointer,
n_vectors, nnz) so repeated per-list calls reuse it.
- Doc sub-matrix A and dense centroid matrix B are built on-device via
gather/scatter kernels reading the resident corpus, so only small id
arrays cross PCIe per inverted list.
- Per-thread cuSPARSE handle + CUDA stream + grow-on-demand scratch
buffers (no global mutex), preserving the 32-way OpenMP concurrency
of build_inverted_lists_clusters().
Assignments are bit-identical to the scalar CPU path (skips centroids,
strict-greater argmax with ties to the lowest cluster index); float
(U32) weights only, with automatic CPU fallback otherwise or on any GPU
error. The hook in map_docs_to_clusters() is gated on
NSPARSE_WITH_CUDA && !__AVX512F__.
Build wiring: NSPARSE_ENABLE_CUDA CMake option (default OFF) plus
cmake/cuda.cmake; CPU-only builds and CI are unaffected.
Tests: tests/gpu_cluster_assigner_test.cpp asserts GPU==CPU cluster
membership (built only with CUDA). All 363 tests pass on an L4.
Benchmark (benchmarks/index_build_benchmark.cpp, Big-ANN SPLADE,
L4 vs 32-core, Release):
- base_small (100K docs, 12.7M nnz): CPU 33.5s -> GPU 6.58s (5.1x)
- base_full (8.84M docs, 1.12B nnz): CPU 483s -> GPU 81.1s (5.96x)
Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
- NSPARSE_GPU_PROFILE=1 env var enables lightweight phase profiling that prints host_prep / gpu_exec / host_assign breakdown at exit. Negligible overhead when disabled (a conditional branch per list). - NSPARSE_BENCH_ITERS / NSPARSE_BENCH_REPS env vars control benchmark iterations/repetitions (default 3/3), so base_full runs can use 1/1 without rebuilding. Profiling findings (base_full, 8.84M docs, L4): gpu_exec dominates at 97% of wall-time (141s aggregate / 84s wall). host_prep and host_assign together are ~5%. GPU SM utilization ~60% during the build; cuSPARSE SpMM throughput is the bottleneck. A custom fused warp-per-row kernel was tested but proved slightly slower than cuSPARSE on base_full's 4400x440 problems (cuSPARSE's tiled algorithm has better B-matrix cache behavior). Kept cuSPARSE. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
The per-term max-pool is ~67% of summarize() and, after the assignment offload, summarize became the dominant remaining CPU cost in build() (profiling: assignment dropped 84x, leaving summarize as the bottleneck that capped GPU utilization at ~45%). Offload the max-pool to the GPU as ONE kernel launch per inverted list (one thread block per cluster), not per cluster. Clusters average only ~10 docs, so a per-cluster launch/sync is dominated by GPU round-trip overhead — an initial per-cluster version was ~5x slower than CPU. The batched per-list kernel collapses ~27M blocking syncs down to ~30K. The tie-order-sensitive sort/truncate stays on the CPU, so output ordering is preserved; results are bit-identical to the CPU max-pool on tie-free data (verified by SummarizeListMaxpool* tests). float (U32) only; default on when a GPU is present, opt out with NSPARSE_GPU_SUMMARIZE=0. Also adds NSPARSE_BUILD_PROFILE / NSPARSE_SUMMARIZE_PROFILE phase timers. Benchmark (base_full, 8.84M docs, L4): build 83.6s -> 74.9s; summarize thread-seconds 1436 -> 760. Overall vs 32-core CPU baseline (483s): 5.96x -> 6.45x. (base_small's ~10-doc clusters are below the break-even point, so GPU summarize is a slight net loss there; the win is at scale.) All 365 tests pass. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
…mmarize Interface cleanup (align with FAISS's GPU model): - Rename CMake option NSPARSE_ENABLE_CUDA -> NSPARSE_ENABLE_GPU and the compile define NSPARSE_WITH_CUDA -> NSPARSE_WITH_GPU. - Remove the NSPARSE_GPU_MIN_DOCS assignment threshold: offload is now unconditional when built with GPU support and a device is present. - Consolidate NSPARSE_BUILD_PROFILE / NSPARSE_SUMMARIZE_PROFILE into the single NSPARSE_GPU_PROFILE env var. - The build benchmark now selects CPU vs GPU at compile time (build with or without NSPARSE_ENABLE_GPU) rather than via a runtime gate. summarize() max-pool: default to a flat dim-sized accumulator on the CPU, replacing the per-cluster std::unordered_map. A per-cluster epoch marks live slots so the dim-sized buffers are reused across clusters with no reset; first-touch is detected via the epoch (not acc==0), robust to zero weights. ~2.8x less CPU work than the map; bit-identical output. GPU summarize becomes opt-in via NSPARSE_GPU_SUMMARIZE=1 (default: CPU flat-array). Measured base_full (8.84M docs, L4): - 32-core: CPU flat-array 63.0s vs GPU summarize 74.9s -> CPU wins - 8-core: CPU flat-array 113.3s vs GPU summarize 101.4s -> GPU wins So the default (CPU) is best on high-core hosts; enable the GPU offload on GPU-rich / low-core hosts where CPU summarize dominates wall-time. All 365 tests pass in both modes (default and NSPARSE_GPU_SUMMARIZE=1). Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
Co-locate the GPU build config with the GPU sources: replace the top-level cmake/cuda.cmake with nsparse/gpu/CMakeLists.txt, which does the CUDA toolchain setup (enable_language, find CUDAToolkit, architectures) and defines a nsparse_attach_gpu(<target>) helper. nsparse/CMakeLists.txt include()s it (when NSPARSE_ENABLE_GPU=ON) after the library targets exist and calls the helper to compile gpu_cluster_assigner.cu into each variant. Also set CUDA_RESOLVE_DEVICE_SYMBOLS ON so device symbols are resolved into the nsparse library itself; pure-C++ consumers (tests, benchmarks) link it without a separate CUDA device-link step. (Previously enable_language ran at top-level scope, which handled this implicitly.) Verified on an L4: GPU build + all 365 tests pass; CPU-only build unaffected. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
chishui
force-pushed
the
gpu-cusparse-build-acceleration
branch
from
July 10, 2026 03:09
e978907 to
6ab574a
Compare
Address review feedback:
- Remove NSPARSE_GPU_PROFILE and all phase-profiling code (Profile struct,
SummarizeProfile, build_prof timers). Profiling will be reintroduced later
in a better form.
- Extract the GPU summarize max-pool into its own GpuSummarizer utility class
(gpu_summarizer.{h,cu}); GpuClusterAssigner now only does assignment. Shared
CUDA helpers and the single resident corpus (GpuCorpus) move to
gpu_common.cuh so both helpers reuse one corpus upload (avoids duplicating
the ~corpus-sized allocation). NSPARSE_GPU_SUMMARIZE=1 still selects GPU vs
CPU summarize; default stays CPU flat-array.
- Tighten comments for concision.
- Document why GPU assignment is excluded from AVX-512 builds: the AVX-512 CPU
path does not skip centroids while the scalar path (which the GPU path
mirrors) does, so mixing them in one build would be inconsistent.
All 365 tests pass in both summarize modes; CPU-only build unaffected.
Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
Replace the interleaved per-cluster GPU checks with a thin dispatcher that picks one path: summarize_with_gpu_ (opt-in, float-only, returns nullopt to fall back) or summarize_with_cpu_. Both feed the shared summarize_emit_cluster_ (sort / alpha-truncate / emit), so output is identical. No behavior change. All 365 tests pass in both summarize modes. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
chishui
marked this pull request as ready for review
July 10, 2026 05:45
chishui
requested review from
model-collapse,
yuye-aws and
zirui-song-18
as code owners
July 10, 2026 05:45
Add GPU build info: the NSPARSE_ENABLE_GPU option, a GPU Acceleration section (what it accelerates, requirements, build commands, the NSPARSE_CUDA_TOOLKIT_ROOT / NSPARSE_CUDA_ARCHITECTURES knobs, and the opt-in NSPARSE_GPU_SUMMARIZE env var), the CUDA/cuSPARSE dependency, and the nsparse/gpu/ directory entry. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Issues Resolved
[#17]
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.