Skip to content

perf(brute_force): add scatter-gather path and fix filtered-search dispatch thresholds#2321

Open
maxwbuckley wants to merge 4 commits into
NVIDIA:mainfrom
maxwbuckley:perf/brute-force-filtered-dispatch-thresholds
Open

perf(brute_force): add scatter-gather path and fix filtered-search dispatch thresholds#2321
maxwbuckley wants to merge 4 commits into
NVIDIA:mainfrom
maxwbuckley:perf/brute-force-filtered-dispatch-thresholds

Conversation

@maxwbuckley

@maxwbuckley maxwbuckley commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

The problem, in one picture

brute_force_search_filtered chooses between a sparse (SDDMM) path and a dense post-filter path on one hardcoded test, sparsity < 0.9f — it runs SDDMM whenever ≤10% of the dataset passes the filter.

But SDDMM's cost climbs with the number of passing rows, while the dense path's is flat (it always does the full GEMM, then masks). The two curves cross long before 10%. cuVS rides the sparse path right past the crossover:

cost of each path

The pink track is what cuVS actually runs. At 10M×128 it is 12.6× slower than necessary at the moment just before it finally gives up on SDDMM.

The green curve is the path this PR adds.

Two things are wrong, not one

1. The threshold is in the wrong place. The true SDDMM/dense crossover is ~3%, not 10%. By the time the current rule hands off, SDDMM is already 2.2–4.4× slower than simply going dense.

2. The threshold is the wrong kind. SDDMM competes against a fixed setup cost, not against work that scales with the dataset, so it stops paying off at an absolute number of passing rows — ~10k, roughly constant across the sweep. Expressed as a fraction, the correct handoff therefore falls as the dataset grows. At 10M rows it is ~0.1%, where the current rule waits until 10%: we keep SDDMM running up to 1,000,000 passing rows when it should hand off at ~10,000. The error grows with the dataset.

Where each path actually wins

regime map

Blue is all that SDDMM should own. Everything left of the red line is what it owns today. The band between them — most of 0.1–10% — belongs to neither of the two paths cuVS currently has.

The change

Three paths, chosen by select_filtered_search_path():

  • sddmm — unchanged; entered only below ~10k passing rows.
  • gather (new) — enumerate the passing rows, compact them (and their precomputed norms) into a dense [n_pass x dim] matrix, run the ordinary unfiltered search over it, map indices back. Cost scales with n_pass instead of n_dataset.
  • dense — unchanged; entered above the gather/dense crossover.

Built from existing primitives — thrust::copy_if over bitset_view::test, raft::matrix::gather, brute_force_knn_impl, thrust::gather. No new kernels.

The gather path applies to bitset filters only, which pass the same rows for every query. A bitmap passes a different set per query, so there is no single set of rows to compact; bitmaps keep a two-way dispatch, but with the corrected ~3% threshold instead of the flat 10%.

What it costs today

cost heatmap

Penalty against an oracle that always picks the fastest of the three paths:

dispatch rule mean p90 worst
sparsity < 0.9 (today) 1.65× 5.29× 12.56×
this PR 1.01× 1.00× 1.49×

Tests

NEIGHBORS_TEST --gtest_filter='*Prefiltered*': 180/180 pass.

Six params were added to selectk_inputs at n_dataset=300000, dim=128, sized to land on each dispatch path. Since these tests check results against a CPU oracle and would pass just as well if the new path never ran, I instrumented the dispatch and confirmed each one is actually reached:

test selectivity passing rows path taken
bitset Result/0 0.10% 300 SDDMM
bitset Result/1–4 7.15% 21,439 GATHER
bitset Result/5 21.0% 63,076 DENSE
bitmap Result/0 0.10% SDDMM
bitmap Result/1–5 7.0–20% DENSE

The four GATHER cases cover L2Expanded, InnerProduct, L2SqrtExpanded and CosineExpanded — the last exercising the gathered-norms path — and all agree with the CPU oracle. Bitmaps correctly carry no passing-row count and never gather.

On the constants

They are empirical, and I want to be plain that they are not derived from first principles. I tried to build a cost model (dense GEMM + gather traffic + select_k) so the dim and n_dataset dependence would fall out of hardware properties; solving it against the measured crossovers produces a negative gather-traffic coefficient, i.e. the model is wrong. So these are fits and I have not dressed them up as anything else.

Instead I kept them round and checked that they tolerate being wrong. Perturbing any constant by 2× — roughly the spread I would expect across GPUs — leaves the mean penalty at 1.01–1.05× and the worst case at or below 2.7×. Still far better than the rule it replaces, which is 12.6× worst case on the hardware it was presumably tuned for. An earlier revision of this branch used 3-significant-figure power-law fits (dim^0.4, dim^0.6, …); they scored identically (1.01× mean), so the precision bought nothing and I removed it.

If maintainers would rather these be tunable via search_params, or auto-calibrated once per device, both are easy from here.

Method: fp16, L2Expanded, k=64, 100 queries, shared bitset; selectivity 0.1–60% × n_dataset ∈ {200K, 500K, 1M, 3M, 10M} × dim ∈ {128, 256, 512, 1024}; 100 reps on an idle RTX 5090 (sm_120), medians (cross-validated against an independent sweep to within 1.6–2.5%).

Known gaps

  1. The thresholds are calibrated for dim >= 128 and n_dataset >= 200K. cuVS runs well below both — this repo's own prefiltered tests use dim from 1 to 2052 at n_dataset=8192. The fits are clamped to the measured range rather than extrapolated, but that corner is unmeasured. All 180 tests pass there; I simply cannot claim the choice is optimal.
  2. One pre-existing test config changes path: n_dataset=8192, dim=8, 10% density moves SDDMM → dense. Both paths are correct (the test passes), but whether that flip is a fix or a regression is exactly what (1) has no data for.
  3. Bitmaps were never benchmarked. The ~3% sparse/dense threshold is carried over from bitset measurements. The cost structure is the same shape, but a bitmap's CSR touches different columns per query, so SDDMM's locality is worse and the true crossover is probably lower than 3% — the error is in the safe direction, but it is inferred, not measured.
  4. Single GPU architecture.
  5. Touches the same function as Promote filtering_rate to base search_params; honor it in brute_force #2120 (filtering_rate → base search_params). No conflict against main today, but the two will need reconciling if both land.

Possible follow-up

A bitmap could get a gather path too, by compacting the union of all queries' passing rows and applying the per-query mask inside that smaller matrix. It degenerates to exactly this PR's path when the filter is a bitset. Whether it pays off depends on how much the queries' filters overlap, so it needs its own benchmark — out of scope here.

…h dispatch thresholds

Filtered brute-force search picks between a sparse (SDDMM) path and a dense
post-filter path on `sparsity < 0.9`, i.e. it runs SDDMM whenever <=10% of the
dataset passes the filter. Two things are wrong with that.

First, the threshold is in the wrong place. SDDMM's cost grows with the number
of passing rows while the dense path's is flat, so across a sweep of
selectivity x n_dataset x dim (fp16, k=64, 100 queries, RTX 5090), SDDMM is
already 2.2-4.4x slower than dense by the time it hands off at 10%. The real
SDDMM/dense crossover is ~2.5-4%.

Second, and more importantly, the threshold is the wrong *kind*. SDDMM stops
paying off at an absolute number of passing rows -- roughly constant in
n_dataset -- not at a fraction of the dataset, because it competes against a
fixed setup cost rather than against work that scales with n_dataset. Expressed
as a fraction, the correct handoff therefore falls as n_dataset grows: at
n_dataset=10M it is ~0.11%, where the current rule waits until 10%. The error
grows with the dataset.

This adds a third path for the band the other two leave uncovered. For a bitset
filter (which selects the same rows for every query) the passing rows are
enumerated, compacted into a dense [n_pass x dim] matrix along with their
precomputed norms, searched with the ordinary unfiltered kernel, and the
resulting indices mapped back. Its cost scales with n_pass rather than
n_dataset, so it wins over most of the 0.1-10% band that SDDMM currently owns.

Measured against the best of the three paths, the current dispatch is up to
12.6x slower than it needs to be (n_dataset=10M, dim=128, 10% selectivity).

A bitmap filter selects different rows per query, so there is no single set of
rows to compact; bitmaps keep a two-way dispatch, but with the corrected
SDDMM/dense threshold rather than the flat 10%.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 14, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

maxwbuckley and others added 3 commits July 14, 2026 21:16
…lerant constants

The previous constants were power-law fits to a single GPU (dim^0.4, dim^0.6,
dim^0.22, 9.3, 212500, 11000). That is false precision: the exponents are not
derivable from a cost model -- solving one against the measured crossovers
yields a negative gather-traffic coefficient -- so they are curve-fits, and
three significant figures of curve-fit do not transfer to other hardware.

Replace them with four round constants. Over the measured sweep this is not a
regression: mean penalty against a perfect oracle is 1.01x either way, and worst
case improves slightly (1.52x -> 1.49x). All the extra precision bought nothing.

The simpler rule also degrades gracefully. Perturbing any constant by 2x, which
is roughly the spread expected across GPUs, keeps the mean penalty at 1.01-1.05x
and the worst case at or below 2.7x -- still far better than the 1.65x mean and
12.6x worst case of the `sparsity < 0.9` rule it replaces.

Also drops the commentary explaining the fits, which no longer exist.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Regret figures against the previous rule describe a code review, not the code.
They tell a future maintainer of this function nothing, and they go stale the
moment anyone retunes. They live in the PR description instead.

Also drops a reference to a benchmark README that is not part of this branch.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A bitmap passes a different set of rows per query, so "the rows it passes" is
not a thing. The dispatch was manufacturing one anyway -- nnz / n_queries -- so
that it could reuse a single variable, then dividing by n_dataset to recover a
selectivity it already had.

Integer division made that lossy: a bitmap passing 5 pairs across 100 queries
gave a row count of 0, hence a selectivity of exactly zero. In general it
understated a bitmap's selectivity and pulled borderline cases onto the sparse
path.

Pass the selectivity (well-defined for both filter kinds, and the quantity the
sparse/dense choice actually turns on) and make the row count a std::optional
that only a bitset supplies -- which is also the only filter that can take the
gather path that needs it.

Dispatch is unchanged except at the sparse/dense boundary for bitmaps, where the
rounding used to be wrong.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@maxwbuckley maxwbuckley marked this pull request as ready for review July 14, 2026 21:28
@maxwbuckley maxwbuckley requested a review from a team as a code owner July 14, 2026 21:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant