cubisect

Hardware-adaptive, GPU-accelerated exact binary search with zero C++ build requirement (NVRTC-only).

Quick start

pip install cubisect

import torch
import cubisect as cb

seq = torch.sort(torch.randn(1024)).values
targets = torch.randn(500_000)

# Dynamic dispatch - autotuned per GPU
idx = cb.dynamic.search(seq, targets)

# Static Eytzinger layout - cache-aware, upload once
handle = cb.static.prepare(seq)
idx = cb.static.search(handle, targets)

Both functions return a torch.int64 tensor of the same shape as targets:

• index of the element if found
• -1 otherwise

Features

• Dynamic - on-device autotune picks the fastest kernel (v0 serial, v2 ILP smem, v2 __ldg global, or an XOR-swizzled smem variant in the CuTe style) based on boundary size and GPU architecture.
• Static - Eytzinger-layout (complete-tree) search in global or shared memory, __ldg-prefetched. Stores only the value tree (4 B/node); the result index is recovered from the search path, so there's no side index array.
• Exact-match semantics - returns the (leftmost) index of each value if present, else -1. Consistent across every kernel, including duplicate values.
• PyTorch-compatible - torch.Tensor in/out, no manual CUDA context.
• Portable - works on any NVIDIA GPU with compute capability ≥ 7.0; compiles kernels at runtime via NVRTC so no C++ toolchain is required.

Hardware autotune

Generate a per-device dispatch table for the dynamic path:

python -m cubisect.autotune \
    --nseq 1,256,1024,8192,65536,131072 \
    --nval 1000000 \
    --warmup 5 --iters 20 --reps 7

Keep --nval large (≥ ~1M): at small query counts the kernels are launch-overhead-bound and the measured differences are pure jitter. Each configuration is timed over --reps repetitions so the table records a median and a standard deviation per kernel.

This produces a JSON file (e.g. dispatch_table_NVIDIA_RTX_4090_sm_89.json), discovered automatically by cb.dynamic.search in this order:

1. the path in $CUBISECT_DISPATCH_TABLE, if set;
2. any dispatch_table*.json shipped next to the installed kernels;
3. any dispatch_table*.json in the current working directory.

The autotuner also includes a margin-of-noise guard: if the best and second-best kernels are within 5 % or within their combined run-to-run standard deviation, the table stores "either" instead of committing to a winner the measurement can't actually distinguish.

Why: CPU vs GPU

This library exists for when you have a sorted array and a large batch of queries, and a threaded CPU binary search is the bottleneck. Binary search is branch-heavy and latency-bound, so the GPU's parallelism is a big win.

Measured on an RTX 5060 Laptop GPU vs a 24-thread CPU (OpenMP), 2^20 queries, exact-match, identical semantics on both sides; median of 7 runs (the plot's shaded bands are min..max). benchmarks/bench_cpu_vs_gpu.py compiles bsearch_cpu.cpp with MSVC and times the GPU with no CPU load running, so the GPU clocks are not thermally throttled.

array size	CPU 1-thread	CPU 24-thread	GPU dynamic	GPU static
256	38 Mq/s	260	16 943	14 707
4 096	24	191	11 546	9 712
65 536	14	124	4 428	6 385
1 048 576	6	78	2 494	3 634
4 194 304	2	36	2 128	3 052

Roughly ~50-60x over a 24-thread CPU and hundreds-to-1000x over a single thread, growing as the array spills CPU cache. There is a crossover near n ~ 16k: dynamic wins for small in-cache arrays (its simpler kernel beats static's rank-accumulation when everything is cached), while static's Eytzinger layout pulls ahead once the array is large enough to be memory-bound. Reproduce with:

python benchmarks/bench_cpu_vs_gpu.py

(Numbers are GPU-specific and meant to show the order of magnitude, not a guarantee. dynamic needs no setup; static adds a one-time prepare() that pays off when you query the same array many times.)

Does the autotuned dispatch help?

cb.dynamic picks a kernel per array size from a per-GPU profile table. The plot below compares that profile-based dispatch against normal dispatch (always using the simple serial v0_global kernel), with each candidate kernel shown as the envelope (benchmarks/bench_dispatch.py):

The dispatch rides the top of the envelope: it picks v2_smem for tiny arrays and the swizzled cute_lite_smem for mid-size arrays (where the tree fits in shared memory), beating always-v0 by ~35-50% for n ≈ 256-4096. For large, out-of-cache arrays every kernel is memory-bound and converges, so v0 is fine and the dispatch correctly falls back to it. Regenerate the table for your GPU with python -m cubisect.autotune (the autotuner queries values spread across the array, so its timings reflect realistic full-depth searches).

Development

pip install -e ".[dev]"
pytest tests/ -v               # full suite (GPU tests auto-skip without CUDA)
pytest tests/test_host.py -v   # CUDA-free unit tests (run anywhere)
python benchmarks/bench_cpu_vs_gpu.py
ruff check src/ tests/
mypy --ignore-missing-imports src/

CI (.github/workflows/ci.yml) runs lint, type-check, and the host tests on GPU-less runners; the full kernel suite needs a self-hosted GPU runner. Pre-commit hooks are configured in .pre-commit-config.yaml.