Add GPU-accelerated IVFPQ search to the Metal backend

[Evandabest]

Summary

This is a follow-up to #5288, which added MetalIndexIVFPQ with GPU-resident storage while delegating search to the CPU.

Adds GPU-accelerated IVFPQ search to the Metal backend, including optimized scanning through M <= 64.

• Add Metal kernels for 8-bit PQ distance computation, inverted-list scanning, and top-k merging
• Support optimized GPU scanning through M <= 64
• Compute distance contributions only for PQ codes encountered in each list when d <= 256
• Add a 32-thread cooperative kernel for lists containing at most 128 vectors
• Retain the cached precomputed-LUT kernel for suitable M <= 16 workloads
• Stream precomputed terms for larger dimensions without allocating an M * 256 threadgroup LUT
• Support both L2 and inner product metrics
• Add exact segmented top-k selection for inverted lists of any length
• Add grouped multi-round merging without a fixed nprobe * k candidate limit
• Compact invalid per-list padding before merging when lists contain fewer than k results
• Reuse dynamically sized Metal scratch buffers across searches
• Retain the legacy per-probe LUT path with optional FP16 lookup tables
• Fall back to CPU search for unsupported configurations, search parameters, or Metal failures

Changes

• Modified: MetalDistance.metal - on demand PQ scans, large-M precomputed scan, short-list specialization, segmented top-k, and compacting grouped merge
• Modified: MetalDistance.h/.mm - IVFPQ search orchestration and support for M <= 64
• Modified: MetalKernels.h/.mm - dispatch for cached, on-demand, short-list, and large-M scan kernels
• Modified: MetalIndexIVFPQ.h/.mm - GPU scan selection, reusable buffers, centroid uploads, precomputed terms, and CPU fallback
• Modified: MetalIndex.h - add the useFloat16 LUT configuration
• Modified: TestMetalIndexIVFPQ.mm - direct-GPU scan validation and search-parameter fallback coverage

Differences from CUDA IVFPQ

Training: MetalIndexIVFPQ::train delegates training to its CPU index, as introduced in #5288. CUDA can train the coarse and product quantizers on GPU. Training is a one-time cost.

Add path: Coarse assignment, residual computation, and PQ encoding remain on the CPU. The encoded PQ codes are stored in GPU-resident Metal buffers. CUDA performs these operations on GPU.

Coarse quantization: The coarse quantizer search runs on the CPU. The selected lists and coarse distances are passed to the Metal scan.

Distance computation: For d <= 256, the preferred Metal path evaluates only the PQ codes encountered while scanning a list. This avoids constructing or loading all 256 lookup entries for every subquantizer, which is especially beneficial for short IVF lists.

For other supported dimensions, Metal uses the CPU IVFPQ precomputed-table decomposition. The query-independent centroid/PQ term is computed once per trained index and the query term once per batch.

Threadgroup memory: The existing cached lookup-table kernel remains available for suitable M <= 16 workloads. Larger M values do not expand its M * 256 threadgroup LUT, which would exceed Metal threadgroup-memory limits. Instead, they use the on-demand or streamed-precomputed scan paths.

List scanning: Short lists use a 32-thread cooperative kernel with approximately 2 KiB of threadgroup scratch. Longer lists use segmented selection to maintain an exact running top-k.

Top-k merge: Per-list results are merged in groups over multiple rounds. Invalid padding is compacted when lists contain fewer than k results.

Fallback: The optimized scan supports M <= 64, d / M <= 256, and k <= 512. Unsupported configurations fall back to the legacy Metal path or CPU search. Search options that the GPU path cannot preserve, including selectors, scan budgets, and polysemous filtering, are forwarded to the CPU with their original parameters.

Build and test

cmake -B build \
  -DFAISS_ENABLE_GPU=OFF \
  -DFAISS_ENABLE_METAL=ON \
  -DBUILD_TESTING=ON \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \
  .

cmake --build build \
  --target faiss faiss_metal TestMetalIndexIVFPQ \
  -j$(sysctl -n hw.logicalcpu)

cd build && ctest -R TestMetalIndexIVFPQ --output-on-failure

All six IVFPQ tests pass. Direct-GPU regression coverage exercises 800 combinations across M=8,16,17,32,48,64, L2 and inner product, k through 512, short and long lists, both merge modes, and 64-bit IDs.

[Evandabest]

Following up on #5288, here are the isolated PQ lookup-table construction results from the GPU-scan work (same as shown in first PR). All measurements were collected on an M3 Pro with d=128, 8-bit PQ, and FP32 lookup tables.

PQ LUT computation
Metal ivfpq_build_lut_l2 kernel vs faiss CPU ProductQuantizer::compute_distance_tables (same M×256 tables, numerically identical, max rel diff ~1e-7):

Case	CPU	Metal	Speedup
M=8, nq=10k	12.3 ms	4.4 ms	2.8x
M=16, nq=10k	17.1 ms	4.6 ms	3.7x
M=32, nq=10k	30.0 ms	8.1 ms	3.7x
M=16, nq=1k	1.78 ms	1.49 ms	1.2x
M=16, nq=50k	84.2 ms	21.9 ms	3.8x
M=16, nq=100k	160.2 ms	43.5 ms	3.7x

LUT speedup saturates at ~3.7-3.8x once the batch is large enough to hide kernel-launch overhead (nq ≥ 10k); ~1.2x for tiny batches.

[meta-codesync]

@mnorris11 has imported this pull request. If you are a Meta employee, you can view this in D113864779.

[mnorris11]

Sorry for the delay. We are a bit behind in PR/issue review. Can you take a look at these when you have time?

• M ≤ 16 ceiling is the main issue. The precomp path caps M at 16 (tgLUT[16*256] = 16KB threadgroup memory, plus 8KB top-k scratch ≈ 24KB of ~32KB). CPU faiss is unbounded, GPU goes to 96, cuVS unbounded, and practical M runs into the hundreds — so most real configs get no GPU speedup. Either tile the LUT over M-blocks so large M runs on GPU, or state the limit explicitly in the summary + header docs.
• Dead code: convert_f32_to_f16 / encodeConvertF32ToF16 are defined and create a pipeline but are never dispatched.
  ivf_scan_list_pq8_precomp early return leaves output uninitialized — the k > PQPRE_MAX_K || M*256 > PQPRE_LUT_MAX guard returns without writing sentinels, unlike the lLen == 0 branch. Unreachable today given host caps, but emits garbage if host/kernel limits drift.
• Collapse ivf_tg_bitonic_int / _long into one template — identical bodies, payload type is the only difference.

[Evandabest]

Hi @mnorris11, no worries. I cleared up the dead code and collapsed ivf_tg_bitonic_int / _long into one template. I am still investigating the M ≤ 16 constraint. It isn't a IVFPQ limitation, it comes from the current kernel's threadgroup memory usage. The lookup table uses 16KB at M=16, and another 8KB for top-k scratch space.

I tested a separate path for M > 16 that reads lookup terms from device memory instead. It works correctly through M=64, but performance is mixed and generally only improves for fine grained indexes. I’m still evaluating whether to include this path or keep and clearly document the existing M <= 16 limit.

I will get back to you by the end of this week

[Evandabest]

Hi @mnorris11, following up on the M <= 16 limitation. I pushed a new optimized Metal IVFPQ scan path supporting M <= 64. Instead of expanding the M * 256 threadgroup LUT, which would exceed Metal’s threadgroup memory limit at larger M, the new path computes distance contributions only for PQ codes actually encountered in each list. Short lists use a 32 thread kernel with about 2 KiB of scratch; longer lists use segmented exact top-k selection. The existing cached LUT kernel remains available for suitable M <= 16 workloads.

[Evandabest]

Benchmark results on an Apple M3 Pro, 18GB RAM are below. Times are milliseconds per 1000 query search, each value is the median of three repeat medians.

d=128, nb=100000, nlist=4096, nprobe=32

M	k	Previous	New	Speedup
8	20	14.98 ms	11.69 ms	1.28x
8	128	18.28 ms	14.24 ms	1.28x
8	512	33.58 ms	19.00 ms	1.77x
16	20	19.66 ms	12.56 ms	1.56x
16	128	20.07 ms	13.35 ms	1.50x
16	512	37.43 ms	17.96 ms	2.08x
32	20	72.10 ms	14.94 ms	4.83x
32	128	34.57 ms	13.70 ms	2.52x
32	512	41.53 ms	19.50 ms	2.13x
64	20	97.06 ms	12.08 ms	8.04x
64	128	54.13 ms	15.15 ms	3.57x
64	512	62.25 ms	23.25 ms	2.68x

Previous represents the old implementation's actual end to end behavior, including CPU fallback where its GPU dispatch limits (M <= 16) were exceeded.