Replication package for the paper Accelerating k-Nearest Neighbors Search with GPU-Based Top-k Selection
Cuda
0
125 commits
updated Sep 24, 2026
This is the replication package for the paper Accelerating k-Nearest Neighbors Search with GPU-Based Top-k Selection submitted to the Journal of Parallel and Distributed Computing (JPDC).
The paper presents a study of state-of-the-art algorithms for top-k selection and distance computation (used in k-NN search) on (NVIDIA) GPUs. It proposes two novel algorithms, bits and bits-fused, outperforming the state-of-the-art algorithms. The proposed algorithms are based on the partial Bitonic sort algorithm by Kruliš et al. (2015) and are enhanced with novel optimizations. All algorithms are implemented in CUDA (CUDA Toolkit) and evaluated on NVIDIA GPUs (Tesla V100, A100, H100, and L40).
To reproduce the results, follow the instructions in Results reproduction. For a quick demonstration of the proposed algorithms and tests validating their correctness, follow the instructions in Demonstration and testing.
Section Plotting the results describes how to visualize the results and maps the produced plots to the figures in the paper. It also discusses the reported speedup over the state-of-the-art algorithms and the relative throughput as a fraction of the theoretical throughput limit deduced from the memory bandwidth of the given GPU.
The replication package contains the following artifacts:
knn, which runs benchmarks for the proposed algorithms and the evaluated state-of-the-art algorithms. The algorithms and their source files are listed in the algorithms.md.*-hopper01-*.csv*-bw01-*.csv*-ampere01-*.csv*-ampere02-*.csv*-volta05-*.csvTo print the manuscript's experimental numbers from the saved runs for all five GPUs, including the latest bw01 runs:
.venv/bin/python scripts/report-manuscript.py
# Save the complete comparison:
.venv/bin/python scripts/report-manuscript.py --show-runs --output plots/manuscript-numbers.md
The script needs NumPy and pandas, and does not run benchmarks or regenerate plots. Edit RUNS, the workload ranges, and the other constants at the top of the script to select different recorded experiments. Defaults use the latest evaluation jobs for each GPU and the saved parameter sweeps for tuning comparisons. The output covers selection, sorting, buffering, prefetching, distance computation, fused kernels, bandwidth utilization, and fixed-parameter slowdowns. Means discard the first ten iterations; the bits tuning runs retain twenty samples and other runs retain ten. --data-dir selects another data directory.
The table below maps the experimental numbers in the manuscript to the row labels printed by report-manuscript.py. Use the H100 (hopper01) column unless another GPU is named. The same rows are reported for every GPU, including bw01. Defaults select the latest runs, so their values can differ from the manuscript's original measurements.
| Manuscript location | Reported number | Script output row / value |
|---|---|---|
| Abstract; §5.2.3, p. 30 | bits 2.1× mean, 1.3× minimum, 2.9× peak speedup | bits speedup vs best competitor: mean / min / max, in that order |
| Abstract; §5.2.3, p. 30 | H100 73% mean, over 87% peak bandwidth utilization | bits bandwidth utilization: mean / peak |
| §5.2.3, p. 30 | A100 79%, V100 89% peak bandwidth utilization | Peak value in bits bandwidth utilization: mean / peak, columns ampere02 and volta05 |
| §5.2.1, p. 28 | Sort-in-registers speedup about 2× for most cases | Sort-in-registers speedup: median; distribution summary in Sort-in-registers speedup: mean / min / max |
| §5.2.1, p. 28 | Fixed block 128, baseline slowdown 31% | Fixed partial-bitonic: block; worst slowdown |
| §5.2.1, p. 28 | Fixed block 128, warp-shuffle slowdown 29% | Fixed partial-bitonic-warp: block; worst slowdown |
| §5.2.1, p. 28 | Fixed block 128, sort-in-registers slowdown 21% | Fixed partial-bitonic-regs: block; worst slowdown |
| §5.2.2, p. 29 | Prefetch gain below 5% overall | Mean value in Prefetch speedup gain: mean / min / max |
| §5.2.3, p. 30 | Mean speedup 1.8× at q=64 | bits speedup at q=64: mean |
| §5.2.3, p. 31; footnote 12 | Fixed bits slowdown 32% H100, 27% A100 | Fixed bits worst slowdown, columns hopper01 and ampere02 |
| Same | H100 block 256, 13 items/thread; A100 block 512, 7 items/thread | Fixed bits parameters: block; items; degree (first component of the items tuple is the batch size) |
| §5.3.1, p. 33 | Distance tuning by d alone: slowdown at most 1.4% | Distance parameters fixed per d: worst slowdown |
| §5.3.2, p. 33 | Extra distance-matrix storage 4 GiB | Two-phase distance matrix (GiB) |
| §5.3.2, p. 34 | bits-fused speedup over 2× versus RAFT for most d<=8 cases | bits-fused vs RAFT, d<=8: speedup >2× gives the count; ratios are in bits-fused vs RAFT, d<=8: mean / min / max |
| §5.3.2, p. 34 | RAFT throughput drops by almost 50% at larger k | RAFT throughput loss at k=64 vs k=8: mean / min / max (the script uses k=8 as the explicit small-k reference) |
| §5.3.2, p. 35 | Fixed bits-fused slowdown 27% H100, 42% A100 | Fixed bits-fused worst slowdown, columns hopper01 and ampere02 |
| §5.3.2, p. 35 | Fixed bits-fused is 6% slower than RAFT in the worst case | Fixed bits-fused vs RAFT: worst overhead (tuning/eval) |
| Same | Worst case d=16, q=1024, k=8 | Fixed bits-fused vs RAFT: worst workload |
| Conclusion, p. 35 | Peak throughput over 75% of the bandwidth limit | Peak value in bits bandwidth utilization: mean / peak |
The manuscript's "all" and "most" comparisons correspond to the rows bits wins vs best competitor, Sort-in-registers fastest, MAGMA-distance wins vs other plotted kernels, bits-fused wins vs RAFT, and bits-fused wins vs two-phase. These counts are winning cases / comparable cases. The per-dimension bits-fused wins vs bits + zero computation rows cover the ideal two-phase comparison in §5.3.2.
Fixed bits-fused versus RAFT comparisons combine the saved tuning and evaluation runs, as indicated by (tuning/eval) in the output.
The following software and hardware are required to reproduce the results (the versions are the ones used in the original experiments, not strict requirements):
cmake version 3.27.6nvcc compiler)
gcc version 13.2.0git version 2.43.5python version 3.9.18 (for visualization)To build the build-NAME/knn binary required to run the benchmarks, run the following commands (replace CUDA_ARCHITECTURES with the compute capability of your GPU, e.g., 90 for NVIDIA H100, and NAME with the chosen name that distinguishes your build from others; the NAME will be used to name the output data as well):
# Clone the repository and update the submodules:
git clone https://github.com/jiriklepl/bits-knn-jpdc2024.git
cd bits-knn-jpdc2024
git submodule update --init --recursive
# Build the `knn` binary in the build-NAME directory:
./local-build.sh NAME CUDA_ARCHITECTURES build
If, for some reason, the ./local-build.sh script does not work, or you want to apply a specific configuration to the build, you can build the source code manually — consult the manual-build.md file for more information. Note that, even when running multiple build jobs in parallel, the build process can take a long time (up to two hours).
All benchmarks are run by run-*.sh scripts in the scripts/ directory. These scripts run the build-NAME/knn binary with various parameters and output the measured data in the CSV format described in data-format.md.
The single runs of the benchmarks are conducted by the build-NAME/knn binary, which is described in more detail in the knn.md file that also describes the available configuration options for the proposed algorithms. For the list of all available parameters and the supported algorithms, run the following command:
build-NAME/knn --help
The local-build.sh script simplifies redirecting the output to separate files data/EXPERIMENT-NAME-TIMESTAMP.csv and data/EXPERIMENT-NAME-TIMESTAMP.err (and also to build and test the source code). This way, we can easily distinguish between various builds (via the NAME parameter) and runs of the benchmarks (TIMESTAMP).
To tune the algorithm parameters, run the following commands (replace CUDA_ARCHITECTURES with the compute capability of your GPU, e.g., 90 for NVIDIA H100, and NAME with the chosen name):
# To benchmark various values for the items-per-thread parameter of the bits algorithm
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-ipt.sh 128 # 128 is the block size
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-ipt.sh 256 # 256 is the block size
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-ipt.sh 512 # 512 is the block size
# To benchmark various parameters of the variants of the `bits-fused` algorithm
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-params.sh
# fused-cache-params is split into three scripts to reduce the running time of each script:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-cache-params.sh 1
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-cache-params.sh 2
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-cache-params.sh 4
# To benchmark various parameters of the optimized Bitonic sort algorithms
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-bitonic-sort.sh
# To benchmark various parameters of WarpSelect and BlockSelect algorithms from FAISS
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-params-warp-select.sh
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-params-block-select.sh
# To benchmark various parameters for various distance computation algorithms
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-distances.sh
This generates various data/EXPERIMENT-NAME-TIMESTAMP.csv and data/EXPERIMENT-NAME-TIMESTAMP.err files. Their format is discussed in data-format.md. Note that these scripts can take a long time (typically up to two hours) to run.
Then, run the following commands to generate scripts/optima.csv and scripts/optima-dist.csv files, which contain the best-performing parameters for the algorithms (they also automatically prepare a virtual environment .venv and install the required Python packages using pip):
# Produce the `scripts/optima.csv` file:
scripts/plot-all.sh analyze-opt
# Produce the `scripts/optima-dist.csv` file:
scripts/plot-all.sh analyze-dist
The two files also produce files scripts/mean-time[-dist]-fixed-max_slowdown.csv and scripts/optima[-dist]-fixed-max_slowdown.csv with the tuning results used to choose one fixed set of parameters for each algorithm (filtered to knn configurations of sufficient query counts) that performs the best across all problem configurations. Each line also shows the slowdown of the run with fixed parameters compared to the best-performing parameters for the given problem configuration (a slowdown of 1.5 would represent a 50% runtime overhead compared to the best-performing run on the given problem configuration). The data in these files are reported in Section 5 in discussions comparing per-configuration tuning to choosing a single fixed parametrization for each algorithm.
To run the benchmarks, run the following commands (assuming you have already built and tuned the algorithms; see the previous sections):
# To run the benchmarks for the optimized Bitonic sort algorithms:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-bitonic-sort.sh
# To run the benchmarks for comparison of the best-performing Bitonic sort algorithm to the proposed bits algorithm (that adds buffering of filtered items):
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-buffer.sh uniform
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-buffer.sh uniform-ascending
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-buffer.sh uniform-descending
# To run the benchmarks comparing the proposed bits algorithm to State-of-the-Art algorithms for top-k selection:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-kselection.sh
# To run the benchmarks for distance computation algorithms:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-distances.sh
# To run the benchmarks for the bits-fused algorithm, comparing it to the bits algorithm and the best-performing distance computation algorithm (MAGMA-distance):
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused.sh
Each command runs the given script and redirects its output to the data/EXPERIMENT-NAME-TIMESTAMP.csv and data/EXPERIMENT-NAME-TIMESTAMP.err files (the format of these files is discussed in data-format.md). Each script should run in a reasonable time (within 30 minutes).
To plot the results, run the following commands:
# (Figure 8 in the paper)
# To plot the results for the optimized Bitonic sort algorithms
scripts/plot-all.sh bitonic-sort
# The same script plots the results for the comparison of the best-performing Bitonic sort algorithm to the proposed bits algorithm:
# (Figure 9 in the paper)
# To merge the results from the uniform, uniform-ascending, and uniform-descending experiments:
scripts/plot-all.sh multibuffer
# (Figure 10 in the paper)
# To plot the results for the proposed bits algorithm compared to State-of-the-Art algorithms for top-k selection:
scripts/plot-all.sh kselection > data/kselection-stats.csv
# (Figure 11 in the paper)
# To plot the results for the distance computation algorithms:
scripts/plot-all.sh distances
# (Figure 12 in the paper)
# To plot the results for the bits-fused algorithm (same as the bits algorithm)
scripts/plot-all.sh kselection
The plots are stored in the plots/ directory; their names normally share the data files' EXPERIMENT-NAME-TIMESTAMP. Multibuffer plots combine only matching identity/ascending/descending runs. The original paper groups are recorded in scripts/multibuffer-runs.json and retain multibuffer-NAME.pdf. New groups require consecutive job IDs in identity/ascending/descending order and identical measurement settings, phases and repetitions; they produce multibuffer-NAME-FIRSTJOB-LASTJOB.pdf. Incomplete or incompatible groups are reported and skipped.
The data/kselection-stats.csv file contains the statistics for the kselection experiment that show the speedup of the proposed algorithm over the state-of-the-art algorithms (maximum, minimum, and average speedup) for each sub-plot and the whole experiment and the relative throughput as a fraction of the throughput limit deduced from the memory bandwidth of the given GPU (the maximum memory throughput divided by the size of one input item in bytes). These reported statistics are used in Section 5.2 of the paper and the Abstract.
Some scripts also produce plots/extra-EXPERIMENT-NAME-TIMESTAMP.pdf files that do not filter out algorithms that have not been chosen for the visualization in the paper (as they were superseded by other algorithms).
The figures in the paper show the following plots:
plots/bitonic-sort-hopper01-108431.pdfplots/multibuffer-hopper01.pdfplots/kselection-hopper01-108424-Uniform.pdfplots/distances-hopper01-108423.pdfplots/fused-hopper01-108416-Uniform.pdfFor more information about the setup used to collect the benchmarking results, see the used-setup.md file.
If you want to quickly demonstrate the proposed algorithms and run the tests validating the correctness of the algorithms, run the following commands (replace CUDA_ARCHITECTURES with the compute capability of your GPU, e.g., 90 for NVIDIA H100, and NAME with the chosen name that distinguishes your build from others):
git clone https://github.com/jiriklepl/bits-knn-jpdc2024.git
cd bits-knn-jpdc2024
git submodule update --init --recursive
# Build the `knn-minimal` binary in the build-NAME directory:
./local-build.sh NAME CUDA_ARCHITECTURES minimal-build
This builds the knn-minimal binary (with reduced configuration options) and the test binary that validates the correctness of the proposed algorithms. On a typical system with sufficient memory (32 GiB or more) and at least 16 CPU cores, the build process should take less than 10 minutes. The supported configuration options are discussed in the knn.md file.
Then, (optionally) run the following commands to run the tests:
./local-build.sh NAME CUDA_ARCHITECTURES test
To demonstrate the efficiency of the proposed algorithms, you can run the following commands:
# Run the proposed bits algorithm with block size 128 and `items-per-thread` set to 16
build-NAME/knn-minimal -a bits -n 1M -q 1k -k 128 --block-size 128 --items-per-thread 16 --seed 42 -r2 | grep "selection"
# Compare it to Air-topk and GridSelect algorithms (state-of-the-art algorithms)
build-NAME/knn-minimal -a air-topk -n 1M -q 1k -k 128 --seed 42 -r2 | grep "selection"
build-NAME/knn-minimal -a grid-select -n 1M -q 1k -k 128 --seed 42 -r2 | grep "selection"
On the standard output, you might see the following lines (the times might vary depending on the hardware):
bits,uniform,identity,0,1048576,1024,1,128,128,"16,1,1",1,selection,2.8317040000000000e-03
bits,uniform,identity,1,1048576,1024,1,128,128,"16,1,1",1,selection,2.6102669999999999e-03
air-topk,uniform,identity,0,1048576,1024,1,128,128,"1,1,1",1,selection,9.3619640000000004e-03
air-topk,uniform,identity,1,1048576,1024,1,128,128,"1,1,1",1,selection,6.6351589999999998e-03
grid-select,uniform,identity,0,1048576,1024,1,128,128,"1,1,1",1,selection,7.9861005999999998e-02
grid-select,uniform,identity,1,1048576,1024,1,128,128,"1,1,1",1,selection,4.4047710000000000e-03
(the outputted Checksums in stderr should be the same for all algorithms; reported times are in seconds - the last column; for a more detailed explanation of the output, see data-format.md).
In the example output above, bits outperforms the state-of-the-art algorithms Air-topk and GridSelect on the given problem size (1M points, 1k queries, 128 nearest neighbors to find) with speedup over $2.5\times$ and $1.7\times$, respectively.
The source code and the replication package are licensed under the MIT License. The license text is available in the LICENSE.txt file.
Some parts of the source code come from third-party projects. Their licenses are available in the licenses.txt file. The relevant parts of the source code are marked with the original license text.
Three applications compare bits, split bits, AIR Top-K, GridSelect and BlockSelect on GPU-resident data:
| Application | Operation | Small / middle / large inputs |
|---|---|---|
| Database top-N | Rank TPC-H rows by discounted price and gather the best N rows | 0.6 / 6 / 60 million rows |
| Token sampling | Select the top-k logits, normalize their probabilities and draw a token | 8 / 128 / 512 sequences, each with 50,257 logits |
| Gradient compression | Select the largest gradient magnitudes and gather their signed values | 0.59 / 2.36 / 38.60 million elements |
The nine cases vary k from 32 to 1024 and sweep bits configurations. Medians use 30 repetitions after 10 warmups, with correctness checks. Full-operator timings cover score preparation, selection and output processing, including launch/synchronization costs. Loading, transfers and model inference/training are excluded.
After the Python setup, run from the repository root, using NAME and CUDA_ARCHITECTURES as above:
./local-build.sh NAME CUDA_ARCHITECTURES applications-prepare
./local-build.sh NAME CUDA_ARCHITECTURES applications-build
./local-build.sh NAME CUDA_ARCHITECTURES applications-run
scripts/plot-all.sh application-scaling
Plotting discovers studies under data/application-scaling/ and writes AIR-relative speedups to plots/application-scaling/. Each study produces 12 combined paper PDFs: three sizes × per-k/global tuning × operator/selection, with all three applications in each file. Every paper PDF has a -configs.csv companion; legends contain only algorithm names. Global plots select one bits variant and fixed configuration independently for each application and size. Detailed PDFs and summary CSVs remain available for each case.
See applications.md for input sizes, measurements, configuration selection and resuming runs.
Replication package for the paper Accelerating k-Nearest Neighbors Search with GPU-Based Top-k Selection
Cuda
0
125 commits
updated Sep 24, 2026
This is the replication package for the paper Accelerating k-Nearest Neighbors Search with GPU-Based Top-k Selection submitted to the Journal of Parallel and Distributed Computing (JPDC).
The paper presents a study of state-of-the-art algorithms for top-k selection and distance computation (used in k-NN search) on (NVIDIA) GPUs. It proposes two novel algorithms, bits and bits-fused, outperforming the state-of-the-art algorithms. The proposed algorithms are based on the partial Bitonic sort algorithm by Kruliš et al. (2015) and are enhanced with novel optimizations. All algorithms are implemented in CUDA (CUDA Toolkit) and evaluated on NVIDIA GPUs (Tesla V100, A100, H100, and L40).
To reproduce the results, follow the instructions in Results reproduction. For a quick demonstration of the proposed algorithms and tests validating their correctness, follow the instructions in Demonstration and testing.
Section Plotting the results describes how to visualize the results and maps the produced plots to the figures in the paper. It also discusses the reported speedup over the state-of-the-art algorithms and the relative throughput as a fraction of the theoretical throughput limit deduced from the memory bandwidth of the given GPU.
The replication package contains the following artifacts:
knn, which runs benchmarks for the proposed algorithms and the evaluated state-of-the-art algorithms. The algorithms and their source files are listed in the algorithms.md.*-hopper01-*.csv*-bw01-*.csv*-ampere01-*.csv*-ampere02-*.csv*-volta05-*.csvTo print the manuscript's experimental numbers from the saved runs for all five GPUs, including the latest bw01 runs:
.venv/bin/python scripts/report-manuscript.py
# Save the complete comparison:
.venv/bin/python scripts/report-manuscript.py --show-runs --output plots/manuscript-numbers.md
The script needs NumPy and pandas, and does not run benchmarks or regenerate plots. Edit RUNS, the workload ranges, and the other constants at the top of the script to select different recorded experiments. Defaults use the latest evaluation jobs for each GPU and the saved parameter sweeps for tuning comparisons. The output covers selection, sorting, buffering, prefetching, distance computation, fused kernels, bandwidth utilization, and fixed-parameter slowdowns. Means discard the first ten iterations; the bits tuning runs retain twenty samples and other runs retain ten. --data-dir selects another data directory.
The table below maps the experimental numbers in the manuscript to the row labels printed by report-manuscript.py. Use the H100 (hopper01) column unless another GPU is named. The same rows are reported for every GPU, including bw01. Defaults select the latest runs, so their values can differ from the manuscript's original measurements.
| Manuscript location | Reported number | Script output row / value |
|---|---|---|
| Abstract; §5.2.3, p. 30 | bits 2.1× mean, 1.3× minimum, 2.9× peak speedup | bits speedup vs best competitor: mean / min / max, in that order |
| Abstract; §5.2.3, p. 30 | H100 73% mean, over 87% peak bandwidth utilization | bits bandwidth utilization: mean / peak |
| §5.2.3, p. 30 | A100 79%, V100 89% peak bandwidth utilization | Peak value in bits bandwidth utilization: mean / peak, columns ampere02 and volta05 |
| §5.2.1, p. 28 | Sort-in-registers speedup about 2× for most cases | Sort-in-registers speedup: median; distribution summary in Sort-in-registers speedup: mean / min / max |
| §5.2.1, p. 28 | Fixed block 128, baseline slowdown 31% | Fixed partial-bitonic: block; worst slowdown |
| §5.2.1, p. 28 | Fixed block 128, warp-shuffle slowdown 29% | Fixed partial-bitonic-warp: block; worst slowdown |
| §5.2.1, p. 28 | Fixed block 128, sort-in-registers slowdown 21% | Fixed partial-bitonic-regs: block; worst slowdown |
| §5.2.2, p. 29 | Prefetch gain below 5% overall | Mean value in Prefetch speedup gain: mean / min / max |
| §5.2.3, p. 30 | Mean speedup 1.8× at q=64 | bits speedup at q=64: mean |
| §5.2.3, p. 31; footnote 12 | Fixed bits slowdown 32% H100, 27% A100 | Fixed bits worst slowdown, columns hopper01 and ampere02 |
| Same | H100 block 256, 13 items/thread; A100 block 512, 7 items/thread | Fixed bits parameters: block; items; degree (first component of the items tuple is the batch size) |
| §5.3.1, p. 33 | Distance tuning by d alone: slowdown at most 1.4% | Distance parameters fixed per d: worst slowdown |
| §5.3.2, p. 33 | Extra distance-matrix storage 4 GiB | Two-phase distance matrix (GiB) |
| §5.3.2, p. 34 | bits-fused speedup over 2× versus RAFT for most d<=8 cases | bits-fused vs RAFT, d<=8: speedup >2× gives the count; ratios are in bits-fused vs RAFT, d<=8: mean / min / max |
| §5.3.2, p. 34 | RAFT throughput drops by almost 50% at larger k | RAFT throughput loss at k=64 vs k=8: mean / min / max (the script uses k=8 as the explicit small-k reference) |
| §5.3.2, p. 35 | Fixed bits-fused slowdown 27% H100, 42% A100 | Fixed bits-fused worst slowdown, columns hopper01 and ampere02 |
| §5.3.2, p. 35 | Fixed bits-fused is 6% slower than RAFT in the worst case | Fixed bits-fused vs RAFT: worst overhead (tuning/eval) |
| Same | Worst case d=16, q=1024, k=8 | Fixed bits-fused vs RAFT: worst workload |
| Conclusion, p. 35 | Peak throughput over 75% of the bandwidth limit | Peak value in bits bandwidth utilization: mean / peak |
The manuscript's "all" and "most" comparisons correspond to the rows bits wins vs best competitor, Sort-in-registers fastest, MAGMA-distance wins vs other plotted kernels, bits-fused wins vs RAFT, and bits-fused wins vs two-phase. These counts are winning cases / comparable cases. The per-dimension bits-fused wins vs bits + zero computation rows cover the ideal two-phase comparison in §5.3.2.
Fixed bits-fused versus RAFT comparisons combine the saved tuning and evaluation runs, as indicated by (tuning/eval) in the output.
The following software and hardware are required to reproduce the results (the versions are the ones used in the original experiments, not strict requirements):
cmake version 3.27.6nvcc compiler)
gcc version 13.2.0git version 2.43.5python version 3.9.18 (for visualization)To build the build-NAME/knn binary required to run the benchmarks, run the following commands (replace CUDA_ARCHITECTURES with the compute capability of your GPU, e.g., 90 for NVIDIA H100, and NAME with the chosen name that distinguishes your build from others; the NAME will be used to name the output data as well):
# Clone the repository and update the submodules:
git clone https://github.com/jiriklepl/bits-knn-jpdc2024.git
cd bits-knn-jpdc2024
git submodule update --init --recursive
# Build the `knn` binary in the build-NAME directory:
./local-build.sh NAME CUDA_ARCHITECTURES build
If, for some reason, the ./local-build.sh script does not work, or you want to apply a specific configuration to the build, you can build the source code manually — consult the manual-build.md file for more information. Note that, even when running multiple build jobs in parallel, the build process can take a long time (up to two hours).
All benchmarks are run by run-*.sh scripts in the scripts/ directory. These scripts run the build-NAME/knn binary with various parameters and output the measured data in the CSV format described in data-format.md.
The single runs of the benchmarks are conducted by the build-NAME/knn binary, which is described in more detail in the knn.md file that also describes the available configuration options for the proposed algorithms. For the list of all available parameters and the supported algorithms, run the following command:
build-NAME/knn --help
The local-build.sh script simplifies redirecting the output to separate files data/EXPERIMENT-NAME-TIMESTAMP.csv and data/EXPERIMENT-NAME-TIMESTAMP.err (and also to build and test the source code). This way, we can easily distinguish between various builds (via the NAME parameter) and runs of the benchmarks (TIMESTAMP).
To tune the algorithm parameters, run the following commands (replace CUDA_ARCHITECTURES with the compute capability of your GPU, e.g., 90 for NVIDIA H100, and NAME with the chosen name):
# To benchmark various values for the items-per-thread parameter of the bits algorithm
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-ipt.sh 128 # 128 is the block size
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-ipt.sh 256 # 256 is the block size
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-ipt.sh 512 # 512 is the block size
# To benchmark various parameters of the variants of the `bits-fused` algorithm
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-params.sh
# fused-cache-params is split into three scripts to reduce the running time of each script:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-cache-params.sh 1
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-cache-params.sh 2
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused-cache-params.sh 4
# To benchmark various parameters of the optimized Bitonic sort algorithms
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-bitonic-sort.sh
# To benchmark various parameters of WarpSelect and BlockSelect algorithms from FAISS
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-params-warp-select.sh
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-params-block-select.sh
# To benchmark various parameters for various distance computation algorithms
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-opt-distances.sh
This generates various data/EXPERIMENT-NAME-TIMESTAMP.csv and data/EXPERIMENT-NAME-TIMESTAMP.err files. Their format is discussed in data-format.md. Note that these scripts can take a long time (typically up to two hours) to run.
Then, run the following commands to generate scripts/optima.csv and scripts/optima-dist.csv files, which contain the best-performing parameters for the algorithms (they also automatically prepare a virtual environment .venv and install the required Python packages using pip):
# Produce the `scripts/optima.csv` file:
scripts/plot-all.sh analyze-opt
# Produce the `scripts/optima-dist.csv` file:
scripts/plot-all.sh analyze-dist
The two files also produce files scripts/mean-time[-dist]-fixed-max_slowdown.csv and scripts/optima[-dist]-fixed-max_slowdown.csv with the tuning results used to choose one fixed set of parameters for each algorithm (filtered to knn configurations of sufficient query counts) that performs the best across all problem configurations. Each line also shows the slowdown of the run with fixed parameters compared to the best-performing parameters for the given problem configuration (a slowdown of 1.5 would represent a 50% runtime overhead compared to the best-performing run on the given problem configuration). The data in these files are reported in Section 5 in discussions comparing per-configuration tuning to choosing a single fixed parametrization for each algorithm.
To run the benchmarks, run the following commands (assuming you have already built and tuned the algorithms; see the previous sections):
# To run the benchmarks for the optimized Bitonic sort algorithms:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-bitonic-sort.sh
# To run the benchmarks for comparison of the best-performing Bitonic sort algorithm to the proposed bits algorithm (that adds buffering of filtered items):
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-buffer.sh uniform
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-buffer.sh uniform-ascending
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-buffer.sh uniform-descending
# To run the benchmarks comparing the proposed bits algorithm to State-of-the-Art algorithms for top-k selection:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-kselection.sh
# To run the benchmarks for distance computation algorithms:
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-distances.sh
# To run the benchmarks for the bits-fused algorithm, comparing it to the bits algorithm and the best-performing distance computation algorithm (MAGMA-distance):
./local-build.sh NAME CUDA_ARCHITECTURES scripts/run-fused.sh
Each command runs the given script and redirects its output to the data/EXPERIMENT-NAME-TIMESTAMP.csv and data/EXPERIMENT-NAME-TIMESTAMP.err files (the format of these files is discussed in data-format.md). Each script should run in a reasonable time (within 30 minutes).
To plot the results, run the following commands:
# (Figure 8 in the paper)
# To plot the results for the optimized Bitonic sort algorithms
scripts/plot-all.sh bitonic-sort
# The same script plots the results for the comparison of the best-performing Bitonic sort algorithm to the proposed bits algorithm:
# (Figure 9 in the paper)
# To merge the results from the uniform, uniform-ascending, and uniform-descending experiments:
scripts/plot-all.sh multibuffer
# (Figure 10 in the paper)
# To plot the results for the proposed bits algorithm compared to State-of-the-Art algorithms for top-k selection:
scripts/plot-all.sh kselection > data/kselection-stats.csv
# (Figure 11 in the paper)
# To plot the results for the distance computation algorithms:
scripts/plot-all.sh distances
# (Figure 12 in the paper)
# To plot the results for the bits-fused algorithm (same as the bits algorithm)
scripts/plot-all.sh kselection
The plots are stored in the plots/ directory; their names normally share the data files' EXPERIMENT-NAME-TIMESTAMP. Multibuffer plots combine only matching identity/ascending/descending runs. The original paper groups are recorded in scripts/multibuffer-runs.json and retain multibuffer-NAME.pdf. New groups require consecutive job IDs in identity/ascending/descending order and identical measurement settings, phases and repetitions; they produce multibuffer-NAME-FIRSTJOB-LASTJOB.pdf. Incomplete or incompatible groups are reported and skipped.
The data/kselection-stats.csv file contains the statistics for the kselection experiment that show the speedup of the proposed algorithm over the state-of-the-art algorithms (maximum, minimum, and average speedup) for each sub-plot and the whole experiment and the relative throughput as a fraction of the throughput limit deduced from the memory bandwidth of the given GPU (the maximum memory throughput divided by the size of one input item in bytes). These reported statistics are used in Section 5.2 of the paper and the Abstract.
Some scripts also produce plots/extra-EXPERIMENT-NAME-TIMESTAMP.pdf files that do not filter out algorithms that have not been chosen for the visualization in the paper (as they were superseded by other algorithms).
The figures in the paper show the following plots:
plots/bitonic-sort-hopper01-108431.pdfplots/multibuffer-hopper01.pdfplots/kselection-hopper01-108424-Uniform.pdfplots/distances-hopper01-108423.pdfplots/fused-hopper01-108416-Uniform.pdfFor more information about the setup used to collect the benchmarking results, see the used-setup.md file.
If you want to quickly demonstrate the proposed algorithms and run the tests validating the correctness of the algorithms, run the following commands (replace CUDA_ARCHITECTURES with the compute capability of your GPU, e.g., 90 for NVIDIA H100, and NAME with the chosen name that distinguishes your build from others):
git clone https://github.com/jiriklepl/bits-knn-jpdc2024.git
cd bits-knn-jpdc2024
git submodule update --init --recursive
# Build the `knn-minimal` binary in the build-NAME directory:
./local-build.sh NAME CUDA_ARCHITECTURES minimal-build
This builds the knn-minimal binary (with reduced configuration options) and the test binary that validates the correctness of the proposed algorithms. On a typical system with sufficient memory (32 GiB or more) and at least 16 CPU cores, the build process should take less than 10 minutes. The supported configuration options are discussed in the knn.md file.
Then, (optionally) run the following commands to run the tests:
./local-build.sh NAME CUDA_ARCHITECTURES test
To demonstrate the efficiency of the proposed algorithms, you can run the following commands:
# Run the proposed bits algorithm with block size 128 and `items-per-thread` set to 16
build-NAME/knn-minimal -a bits -n 1M -q 1k -k 128 --block-size 128 --items-per-thread 16 --seed 42 -r2 | grep "selection"
# Compare it to Air-topk and GridSelect algorithms (state-of-the-art algorithms)
build-NAME/knn-minimal -a air-topk -n 1M -q 1k -k 128 --seed 42 -r2 | grep "selection"
build-NAME/knn-minimal -a grid-select -n 1M -q 1k -k 128 --seed 42 -r2 | grep "selection"
On the standard output, you might see the following lines (the times might vary depending on the hardware):
bits,uniform,identity,0,1048576,1024,1,128,128,"16,1,1",1,selection,2.8317040000000000e-03
bits,uniform,identity,1,1048576,1024,1,128,128,"16,1,1",1,selection,2.6102669999999999e-03
air-topk,uniform,identity,0,1048576,1024,1,128,128,"1,1,1",1,selection,9.3619640000000004e-03
air-topk,uniform,identity,1,1048576,1024,1,128,128,"1,1,1",1,selection,6.6351589999999998e-03
grid-select,uniform,identity,0,1048576,1024,1,128,128,"1,1,1",1,selection,7.9861005999999998e-02
grid-select,uniform,identity,1,1048576,1024,1,128,128,"1,1,1",1,selection,4.4047710000000000e-03
(the outputted Checksums in stderr should be the same for all algorithms; reported times are in seconds - the last column; for a more detailed explanation of the output, see data-format.md).
In the example output above, bits outperforms the state-of-the-art algorithms Air-topk and GridSelect on the given problem size (1M points, 1k queries, 128 nearest neighbors to find) with speedup over $2.5\times$ and $1.7\times$, respectively.
The source code and the replication package are licensed under the MIT License. The license text is available in the LICENSE.txt file.
Some parts of the source code come from third-party projects. Their licenses are available in the licenses.txt file. The relevant parts of the source code are marked with the original license text.
Three applications compare bits, split bits, AIR Top-K, GridSelect and BlockSelect on GPU-resident data:
| Application | Operation | Small / middle / large inputs |
|---|---|---|
| Database top-N | Rank TPC-H rows by discounted price and gather the best N rows | 0.6 / 6 / 60 million rows |
| Token sampling | Select the top-k logits, normalize their probabilities and draw a token | 8 / 128 / 512 sequences, each with 50,257 logits |
| Gradient compression | Select the largest gradient magnitudes and gather their signed values | 0.59 / 2.36 / 38.60 million elements |
The nine cases vary k from 32 to 1024 and sweep bits configurations. Medians use 30 repetitions after 10 warmups, with correctness checks. Full-operator timings cover score preparation, selection and output processing, including launch/synchronization costs. Loading, transfers and model inference/training are excluded.
After the Python setup, run from the repository root, using NAME and CUDA_ARCHITECTURES as above:
./local-build.sh NAME CUDA_ARCHITECTURES applications-prepare
./local-build.sh NAME CUDA_ARCHITECTURES applications-build
./local-build.sh NAME CUDA_ARCHITECTURES applications-run
scripts/plot-all.sh application-scaling
Plotting discovers studies under data/application-scaling/ and writes AIR-relative speedups to plots/application-scaling/. Each study produces 12 combined paper PDFs: three sizes × per-k/global tuning × operator/selection, with all three applications in each file. Every paper PDF has a -configs.csv companion; legends contain only algorithm names. Global plots select one bits variant and fixed configuration independently for each application and size. Detailed PDFs and summary CSVs remain available for each case.
See applications.md for input sizes, measurements, configuration selection and resuming runs.