Skip to content

Score expanded-L2 brute-force tiles with RAFT's GEMM bias epilogue - #3049

Draft
tarang-jain wants to merge 7 commits into
NVIDIA:mainfrom
tarang-jain:bf-l2-gemm-bias
Draft

Score expanded-L2 brute-force tiles with RAFT's GEMM bias epilogue#3049
tarang-jain wants to merge 7 commits into
NVIDIA:mainfrom
tarang-jain:bf-l2-gemm-bias

Conversation

@tarang-jain

@tarang-jain tarang-jain commented Nov 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Note

Description

For L2Expanded and L2SqrtExpanded, tiled_brute_force_knn computes each tile in two steps. An inner-product GEMM writes q.x, then a map_offset pass reads and rewrites the whole tile to turn each value into |q|^2 + |x|^2 - 2 q.x. Only then does select_k run. This PR removes that pass:

  • One raft::linalg::gemm call per tile writes |x|^2 - 2 q.x directly. The call computes search * index^T with alpha = -2. Its new bias argument (Add an optional bias epilogue to raft::linalg::gemm raft#3205, cuBLASLt's bias epilogue) adds the FP32 index norms, one per index row, inside the GEMM.
  • select_k ranks those scores as they are, because |q|^2 is the same along a row. Afterwards, only the k values kept per row get |q|^2 added, clamped at 0 and square-rooted, the same arithmetic as l2_exp_cutlass_op.
  • Arithmetic is unchanged. Float stays plain FP32 (CUBLAS_COMPUTE_32F, no TF32). Half runs on the tensor cores with FP32 accumulation, as before.
  • Scope. The new path applies to float or half rows with FP32 distances and no custom distance epilogue. Cosine, other metrics and custom epilogues are unchanged.
  • Alignment. cuBLASLt's fast kernels need the bias and the tile's rows 16-byte aligned. When the index is split into column tiles, each tile spans a multiple of 4 index rows (and at least k). Precomputed index norms that aren't 16-byte aligned, such as a slice of a larger norm vector, are copied once per search into an aligned buffer. Add an optional bias epilogue to raft::linalg::gemm raft#3205 chooses the GEMM's algorithm for the bias's actual alignment, so a misaligned bias would still be correct, just slower.

Benchmarks

brute_force::build + brute_force::search with L2Expanded on an H100 80GB HBM3. The baseline is #3041 (a3694728e) and the comparison is this branch. Both were built from source for sm_90, with this branch built against NVIDIA/raft#3205's headers. Shapes: half and float, dim 96 and 768, 1K–256K index rows, 16/256/4096 queries, k = 10 and 100. Each time is the mean over every timed call: 3 batches of equal length, about 300 ms each.

  • Half, 64K–256K rows, 256–4096 queries: 1.16–1.66× faster (mean 1.41×). For example, dim 96, 256K rows, 4096 queries, k = 10: 7.36 → 4.43 ms.
  • Half, 8K rows, 4096 queries: 1.23–1.46× faster (mean 1.33×).
  • Half, smaller shapes (16 queries, or 1K–8K rows with up to 256 queries, which take about 20–80 µs per search): 1.02–1.25× (mean 1.07×).
  • Float, k = 100:
    • 4096 queries: 1.05–1.28× faster (mean 1.13×).
    • 256 queries: 1.02–1.26× faster (mean 1.09×).
    • 16 queries: 1.01–1.06× (mean 1.03×).
  • Float, k = 10: unchanged (0.99–1.01×), because float with k <= 64 still uses fusedL2Knn.

Alignment, same setup, half, k = 10. "Before" is this PR and NVIDIA/raft#3205 without their alignment commits:

  • Norms offset by one float, as a slice of a larger norm vector is: before, search failed with "misaligned address" at 128 and 1024 queries over 20K rows, at dim 96 and 768. Now it runs. Copying the norms costs 1–4 µs per search over aligned norms.
  • 500 queries over 300K rows, where column tiles were 268,435 rows wide: dim 768 went from 2.92 to 0.91 ms, dim 96 from 0.91 to 0.70 ms. 100 queries over 3M rows, dim 96: 1.93 → 1.48 ms.
  • Aligned norms, tiles already a multiple of 4 rows: unchanged.

Results match #3041: distances agree to within 2.2e-07 relative at every rank. Neighbor ids differ only between equal-distance ties; the lowest recall against #3041 is 0.99996.

Full table (96 shapes)
dtype dim index rows queries k #3041 (ms) this branch (ms) speedup
half 96 1K 16 10 0.0184 0.0176 1.05×
half 96 1K 16 100 0.0191 0.0185 1.03×
half 96 1K 256 10 0.0225 0.0211 1.07×
half 96 1K 256 100 0.0241 0.0235 1.03×
half 96 1K 4096 10 0.0491 0.0392 1.25×
half 96 1K 4096 100 0.0836 0.0742 1.13×
half 96 8K 16 10 0.0186 0.0179 1.04×
half 96 8K 16 100 0.0304 0.0295 1.03×
half 96 8K 256 10 0.0425 0.0376 1.13×
half 96 8K 256 100 0.0759 0.0709 1.07×
half 96 8K 4096 10 0.2941 0.2010 1.46×
half 96 8K 4096 100 0.3658 0.2863 1.28×
half 96 64K 16 10 0.0369 0.0339 1.09×
half 96 64K 16 100 0.0636 0.0606 1.05×
half 96 64K 256 10 0.1472 0.1070 1.38×
half 96 64K 256 100 0.2230 0.1806 1.23×
half 96 64K 4096 10 1.8656 1.1649 1.60×
half 96 64K 4096 100 2.5410 1.8601 1.37×
half 96 256K 16 10 0.0772 0.0679 1.14×
half 96 256K 16 100 0.1317 0.1218 1.08×
half 96 256K 256 10 0.4827 0.3089 1.56×
half 96 256K 256 100 0.6400 0.4626 1.38×
half 96 256K 4096 10 7.3586 4.4265 1.66×
half 96 256K 4096 100 9.5199 6.5541 1.45×
half 768 1K 16 10 0.0182 0.0175 1.04×
half 768 1K 16 100 0.0202 0.0197 1.03×
half 768 1K 256 10 0.0234 0.0223 1.05×
half 768 1K 256 100 0.0251 0.0246 1.02×
half 768 1K 4096 10 0.0552 0.0465 1.19×
half 768 1K 4096 100 0.0902 0.0812 1.11×
half 768 8K 16 10 0.0202 0.0186 1.09×
half 768 8K 16 100 0.0321 0.0311 1.03×
half 768 8K 256 10 0.0455 0.0406 1.12×
half 768 8K 256 100 0.0785 0.0742 1.06×
half 768 8K 4096 10 0.3406 0.2525 1.35×
half 768 8K 4096 100 0.4145 0.3374 1.23×
half 768 64K 16 10 0.0738 0.0710 1.04×
half 768 64K 16 100 0.0996 0.0973 1.02×
half 768 64K 256 10 0.1840 0.1445 1.27×
half 768 64K 256 100 0.2599 0.2231 1.16×
half 768 64K 4096 10 2.2572 1.5316 1.47×
half 768 64K 4096 100 2.9045 2.2198 1.31×
half 768 256K 16 10 0.2097 0.1961 1.07×
half 768 256K 16 100 0.2663 0.2529 1.05×
half 768 256K 256 10 0.6254 0.4415 1.42×
half 768 256K 256 100 0.7919 0.5958 1.33×
half 768 256K 4096 10 9.1236 6.0714 1.50×
half 768 256K 4096 100 11.0350 7.9752 1.38×
float 96 1K 16 10 0.0480 0.0476 1.01×
float 96 1K 16 100 0.0197 0.0187 1.05×
float 96 1K 256 10 0.0491 0.0489 1.00×
float 96 1K 256 100 0.0261 0.0256 1.02×
float 96 1K 4096 10 0.1593 0.1597 1.00×
float 96 1K 4096 100 0.1005 0.0909 1.11×
float 96 8K 16 10 0.1320 0.1330 0.99×
float 96 8K 16 100 0.0326 0.0317 1.03×
float 96 8K 256 10 0.1245 0.1253 0.99×
float 96 8K 256 100 0.0849 0.0798 1.06×
float 96 8K 4096 10 0.5624 0.5622 1.00×
float 96 8K 4096 100 0.4998 0.4205 1.19×
float 96 64K 16 10 0.7022 0.7003 1.00×
float 96 64K 16 100 0.0778 0.0753 1.03×
float 96 64K 256 10 0.4138 0.4136 1.00×
float 96 64K 256 100 0.2989 0.2583 1.16×
float 96 64K 4096 10 3.6755 3.6909 1.00×
float 96 64K 4096 100 3.5320 2.8617 1.23×
float 96 256K 16 10 0.7730 0.7716 1.00×
float 96 256K 16 100 0.1826 0.1728 1.06×
float 96 256K 256 10 1.4935 1.4929 1.00×
float 96 256K 256 100 0.8939 0.7100 1.26×
float 96 256K 4096 10 14.4611 14.4608 1.00×
float 96 256K 4096 100 13.3068 10.3560 1.28×
float 768 1K 16 10 0.0976 0.0980 1.00×
float 768 1K 16 100 0.0233 0.0222 1.05×
float 768 1K 256 10 0.0987 0.0987 1.00×
float 768 1K 256 100 0.0417 0.0408 1.02×
float 768 1K 4096 10 0.4823 0.4834 1.00×
float 768 1K 4096 100 0.2133 0.2039 1.05×
float 768 8K 16 10 0.1823 0.1836 0.99×
float 768 8K 16 100 0.0576 0.0567 1.02×
float 768 8K 256 10 0.3226 0.3203 1.01×
float 768 8K 256 100 0.1451 0.1379 1.05×
float 768 8K 4096 10 3.3799 3.3534 1.01×
float 768 8K 4096 100 1.3604 1.2723 1.07×
float 768 64K 16 10 0.8364 0.8336 1.00×
float 768 64K 16 100 0.1641 0.1622 1.01×
float 768 64K 256 10 2.6213 2.6244 1.00×
float 768 64K 256 100 0.7196 0.6837 1.05×
float 768 64K 4096 10 27.0711 27.0758 1.00×
float 768 64K 4096 100 10.3987 9.7342 1.07×
float 768 256K 16 10 1.2001 1.1959 1.00×
float 768 256K 16 100 0.5003 0.4926 1.02×
float 768 256K 256 10 10.3261 10.3217 1.00×
float 768 256K 256 100 2.5980 2.4124 1.08×
float 768 256K 4096 10 108.4281 108.4351 1.00×
float 768 256K 4096 100 40.7220 37.6692 1.08×

Testing

All of the following were built against NVIDIA/raft#3205's headers.

  • NEIGHBORS_ANN_BRUTE_FORCE_TEST: 78/78 pass.
  • NEIGHBORS_TEST: 373/373 pass. This covers the prefiltered bitmap and bitset tests with k = 0, the half and float tests, and odd dims.
  • pre-commit run --all-files passes.

🤖 Generated with Claude Code

tarang-jain and others added 5 commits November 10, 2026 22:57
@tarang-jain @claude
Use tiled GEMM brute force for half with k <= 64 and grow tile rows f…
837208a 
…or small indexes

fusedL2Knn has no tensor-core path, so half inputs with k <= 64 now take
tiled_brute_force_knn. Float stays on the fused kernel.

When the whole index fits in one column tile, chooseTileSize now grows
tileRows up to the existing memory target (capped at numQueries) instead of
stopping at 512, so small indexes need fewer GEMM + select_k rounds. Large
indexes are unchanged. The per-column-tile merge buffers are only allocated
when tiling over columns.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@tarang-jain @claude
Clarify when the brute-force merge buffers are used
182cf17 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@tarang-jain
Merge branch 'main' into bf-half-tiled-knn
4cb9d8b
@tarang-jain @claude
Size brute-force tiles against k when the index has fewer than k rows
a369472 
tiled_brute_force_knn raises tile_cols to k after chooseTileSize, so with
n < k the row tile sized against n made the distance buffer far larger than
the memory target (n=10, k=1000, m=1e6: 4 GB). Pass max(n, k) as the column
count so rows are sized against the actual tile width. Unchanged for k <= n.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@tarang-jain @claude
Score expanded-L2 brute-force tiles with a cuBLASLt bias epilogue
27cf1e3 
For L2Expanded and L2SqrtExpanded over float or half rows (FP32 distances, no custom
epilogue), each tile is now one cuBLASLt matmul that writes |x|^2 - 2 q.x directly:
alpha = -2 scales the dot products and the bias epilogue adds the FP32 index norms. This
replaces the inner-product GEMM plus the pass that turned every dot product into a
distance. select_k ranks the scores as they are, since |q|^2 is constant along a row, and
only the k kept per row get |q|^2 added, clamped and rooted as l2_exp_cutlass_op does.

The matmul is planned per tile shape and per operand alignment (tiles start at arbitrary
rows), cached through a custom resource, and falls back to the previous path if cuBLASLt
has no algorithm for it. Float stays plain FP32 (no TF32); half runs on the tensor cores.

H100, brute_force::search vs the parent commit: half 1.3-1.66x faster at 256-4096 queries
over 8k-262k rows, float with k = 100 up to 1.28x, results identical up to ties.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@tarang-jain tarang-jain added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels Nov 12, 2026
@copy-pr-bot

copy-pr-bot Bot commented Nov 12, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@tarang-jain tarang-jain mentioned this pull request Nov 12, 2026
Draft
@tarang-jain @claude
Use raft::linalg::gemm's bias epilogue instead of calling cuBLASLt
df06c95 
The tile's |x|^2 - 2 q.x now comes from raft::linalg::gemm (column-major: index * search^T,
alpha = -2) with the index norms as its bias, which NVIDIA/raft#3205 adds. This drops
l2_expanded_scores.cuh and its direct cuBLASLt calls, descriptor cache and alignment
handling. Same arithmetic, tests and timings as before.

Requires NVIDIA/raft#3205.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@tarang-jain tarang-jain changed the title Score expanded-L2 brute-force tiles with a cuBLASLt bias epilogue Score expanded-L2 brute-force tiles with RAFT's GEMM bias epilogue Nov 12, 2026
@tarang-jain @claude
Keep the L2 GEMM's bias and tile rows 16-byte aligned
5634bff 
cuBLASLt runs slower kernels when the bias (the index norms at norms + j) or the rows of
the distance tile aren't 16-byte aligned. When the index is split into column tiles, each
tile now spans a multiple of 4 index rows, and precomputed index norms that aren't
16-byte aligned, such as a slice of a larger norm vector, are copied once per search into
an aligned buffer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

Unstructured Data Processing
Status: No status

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@tarang-jain