feat: GPU HNSW search with parallel beam-search kernel — native FP32/FP16/BF16/INT8

[devin-ai-integration]

Summary

Adds GpuIndexHNSW — a GPU implementation of HNSW search using a custom parallel beam-search kernel. Accepts a CPU-built IndexHNSW/IndexHNSWSQ via copyFrom() / copyFromWithMetric(), uploads vectors + graph to VRAM, frees the CPU copy, and runs search entirely on GPU. Recall is tuned via ef.

The device dataset is stored in four native precisions, each kept in its on-disk layout (no up-conversion at upload) and up-converted to fp32 per element inside the kernel for distance accumulation:

type	source storage	device bytes/elem	kernel load_elem
FP32	HNSW-Flat	4	float
FP16	HNSW_SQ QT_fp16	2	__half2float
BF16	HNSW_SQ QT_bf16	2	__bfloat162float
INT8	HNSW_SQ QT_8bit_direct_signed	1	static_cast<float>

Dispatch is driven by an explicit enum instead of a single bool flag:

enum class GpuHnswDatasetType { FP32, INT8, FP16, BF16 };
// GpuHnswSearch.cuh
switch (idx.dataset_type) {
  case FP16: launch_kernels(static_cast<const half*>(idx.d_dataset), ...); break;
  case BF16: launch_kernels(static_cast<const __nv_bfloat16*>(idx.d_dataset), ...); break;
  case INT8: launch_kernels(static_cast<const int8_t*>(idx.d_dataset), ...); break;
  default:   launch_kernels(static_cast<const float*>(idx.d_dataset), ...); break;
}

FP16/BF16 mirror the pre-existing INT8 native path: from_faiss_hnsw_sq() copies the raw 2-byte SQ codes verbatim via upload_halfwidth_dataset() (bit-compatible with CUDA half/__nv_bfloat16), and for COSINE computes d_inv_norms from the fp32-decoded values (stored codes are not normalized). The templated distance helpers (thread_l2_distance/thread_ip_distance) were already DataT-generic; only the new load_elem overloads and the dispatch were added. The graph-walk algorithm is unchanged.

Key design decisions:

1. Parallel beam-search kernel (GpuHnswSearchKernel.cuh): Layer-0 search uses a sorted result buffer (ef entries in smem) + a staging buffer (sw × max_degree0) merged by thread 0 via insertion sort. Stagnation is detected via num_parents == 0 (early exit). An earlier "OCQ" overflow queue was removed as dead code — see the decision record docs/design-docs/design_docs/20260711-gpu-hnsw-ocq-removal.md in 6si/milvus#6. Use a larger ef for higher recall.

2. 1-thread-per-distance: Each of 128 threads computes one full distance independently — 128 concurrent distances vs 32 with 4-thread cooperative; the concurrency outweighs coalescing at these vector sizes.

3. Native low-precision storage: FP16/BF16 (2 bytes) and INT8 (1 byte) stay in their native layout on device — 2×/4× less VRAM than fp32 (critical at tens of millions of rows/node) — with per-element up-conversion to fp32 in the kernel.

4. Per-segment streams + scratch pool: each GpuHnswDeviceIndex owns a cudaStreamNonBlocking; searches acquire a GpuHnswSearchScratch slot (+ stream) from GpuHnswScratchPool for concurrent execution across Milvus segments.

Host API:

gpu_index->searchHost(nq, h_queries, k, h_distances, h_labels, params);
// params.ef, params.search_width, params.max_iterations, params.thread_block_size

Verification: compile-verified via the GPU build; GPU-runtime recall/latency/VRAM validated on-cluster — see Performance below and the unit/recall/sanitizer gates (fp32/fp16/bf16/int8 × L2/IP/COSINE).

Performance (CPU-HNSW vs GPU)

Benchmarked on the same dataset and harness (collection mpd_v2, 569.9M rows, INT8 COSINE; run_eval_jobs.py, --sample-size 2048 --continuous-iterations 30). Best cell = batch 512, 8 parallel × 4 workers (32 concurrency):

Cluster	Nodes	Throughput	p50 latency	R@1
GPU HNSW (g7e.2xlarge)	8	28,087 vec/s	518 ms	0.897
CPU HNSW (r8g.2xlarge Graviton)	32	4,365 vec/s	3,683 ms	0.889
CPU HNSW (r8g.2xlarge Graviton)	10	1,509 vec/s	9,178 ms	0.891

• GPU on 8 nodes ≈ 6.4× the throughput of CPU on 32 nodes and ≈ 18.6× the 10-node CPU cluster, at ~7× lower p50 latency.
• Peak GPU throughput = 30,913 vec/s (batch 256, 8 parallel × 1 worker; p50 62 ms, R@1 0.897).
• Recall parity throughout: R@1 ≈ 0.889–0.900 on both CPU and GPU.
• CPU HNSW is CPU-bound — each node pins at its 8-vCPU limit; throughput scales ~linearly with node count while extra concurrency only inflates latency (CPU 8×4 b1024 → p50 18s, p99 30s). GPU sweet spot is batch 256–512; batch 1024 only helps at the highest concurrency.

Throughput figures are from the gpu-hnsw-v89 image (this beam-search kernel; fp32/int8 COSINE). v89 is 10–24% faster than v83 at all matched batch-512 cells (4×1 +11%, 4×4 +11%, 8×1 +24%, 8×4 +14%), so the filtered-search path does not regress unfiltered throughput. Native fp16/bf16 throughput is a follow-up; correctness across fp32/fp16/bf16/int8 × L2/IP/COSINE (and filtered deletes/TTL/partition at CPU parity) is covered by the unit/recall/sanitizer gates and the sealed-segment cluster e2e.

Related PRs: 6si/knowhere#2 (Knowhere wrapper + registrations), 6si/milvus#6 (Milvus integration).

Link to Devin session: https://6sense.devinenterprise.com/sessions/d39aba56b8a3467cbbf231ab631a06ed

Link to Devin session: https://6sense.devinenterprise.com/sessions/55dab0bd33c346df99b5818b30059342
Requested by: @premal

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[devin-ai-integration]

Pushed 1638541c7 addressing an external code review of the GPU HNSW code:

• H1 — GpuHnswScratchPool::release() now cudaStreamSynchronizes the slot's stream before returning it to the pool. Callers (searchHost/searchImpl_) already sync, so this is normally a no-op; it's a safety net against a future caller that skips its own sync (whose in-flight buffers could otherwise be freed+realloced by the next acquirer's ensure()).
• L2 — GpuIndexHNSW docstring updated to list fp16/bf16 (native 2-byte device storage) alongside float32/int8.
• C1 — documented that searchHost is the preferred entry point and searchImpl_ round-trips labels D2H->H2D (a GPU-side label-conversion kernel is a follow-up; Knowhere uses searchHost).
• H5 — documented the HnswT levels[] 1-based-count convention required by extract_hnsw_layers.
• L1 — documented num_layers as informational (search uses num_upper_layers_built).

Other review items were verified as already-correct or intentional (const_cast, cudaSetDevice-in-dtor, ef-not-validated matching the CPU checker). The C1 GPU-kernel label-conversion and the cosine-recall investigation (H4) are tracked as follow-ups; GPU unit tests are running now on an L40S.

[devin-ai-integration]

Codex P1 #1 — upload CPU inverse L2 norms for quantized cosine (2f5150915)

The knowhere CPU cosine index records inverse L2 norms from the original input vectors (HasInverseL2Norms::get_inverse_l2_norms()). This GPU build previously recomputed norms from decoded/quantized codes, and the fp32 fallback normalized decoded vectors in place — both diverge from the CPU index's scores and HNSW traversal for lossy SQ8/int8/fp16/bf16 cosine.

Fix: thread the stored inverse-norm pointer from from_faiss_hnsw_sq into the int8 / halfwidth (fp16/bf16) / fp32 upload paths and apply it in the kernel. Only flat-fp32 cosine (decoded == original) keeps normalizing in place. throw if a cosine SQ index is missing its inverse norms.

Byte-identical with the knowhere vendored copy (6si/knowhere#2 08ced662); diff is clean. Runtime validation via the new knowhere [gpu_hnsw_p1] cosine-parity test on L40S (v67, queued).

[premal]

Code Review — Delta (DP4A + Padded Staging)

⚠️ Still Open from Previous Review

[I8] No validation that k ≤ ef — If a caller sets k > ef, the kernel silently returns fewer than k results (padded with sentinels UINT64_MAX / FLT_MAX). Consider adding FAISS_THROW_IF_NOT_MSG(k <= ef, ...) in searchHost / searchImpl_, or auto-adjusting ef = max(ef, k).

New Changes — No New Critical Issues

DP4A path is correct. __dp4a usage with aligned packed int32, dim%4==0 guard at dispatch, no int32 overflow risk. Bit-exact for QT_8bit_direct_signed. The QueryT/USE_DP4A template threading via if constexpr is clean.

Minor suggestions on new code:

• [D1] Add a comment on the DP4A dispatch explaining that use_ip is the gate because DP4A only computes dot products (not L2). Currently the guard sc.d_queries_i8 != nullptr && idx.use_ip && (dim % 4 == 0) is correct but the rationale isn′t documented.

• [D2] searchHostInt8 does a full cudaStreamSynchronize after uploading fp32_q (for host buffer lifetime safety). This serializes the stream before the kernel launches — consider pinned memory + deferred sync as a follow-up.

Previously flagged, still present (low priority):

• #include <atomic> in GpuIndexHNSW.h is unused
• GpuHnswScratchSlot destructor missing cudaSetDevice before cudaStreamDestroy (multi-GPU safety)
• cudaFuncSetAttribute called on every search (correct but wasteful; should cache per template instantiation)

Positives: Padded staging is consistently used across kernel, smem calc, and merge call sites. The d_layer0_queries vs d_queries distinction correctly keeps upper-layer greedy descent on fp32 while layer-0 uses int8. The ef clamp warning and block-size error messages are good UX improvements.

Full review: static analysis only (no CUDA/GPU available for compile/test verification).

[premal]

Code Review — Delta 2 (Bitonic Sort Race Fix + Defensive Guards)

Both bug fixes verified correct via brute-force simulation

1. Bitonic sort race fix — ✅ Correct
The two-phase barrier-separated compare-exchange (read-all → __syncthreads → write-all → __syncthreads) eliminates the cross-warp shared-memory hazard. Verified via:

• Structural proof: do_swap is always symmetric for both threads in a pair
• 50,000 full bitonic sort trials: 0 mismatches vs reference sort
• 100,000 full merge trials: 0 mismatches, 0 sentinel promotions

2. Merge-path ri_mid off-by-one fix — ✅ Critical bug fixed
The old formula ri_mid = pos - mid was wrong. Brute-force over 1.65M test cases:

• Old formula: 50.02% mis-partitioned — explains the ~4% recall loss and sentinel-promotion crashes
• New formula (pos - mid - 1): 0% wrong across all edge cases

This was a real correctness bug, not just a defensive guard.

3. Defensive guards — ✅ Correct, don't mask live bugs
With the merge fix in place, sentinels never enter result_ids[0..rc) on the healthy path (0 sentinel promotions in 100K trials). The guards are belt-and-suspenders against OOB faults — good practice.

Still Open (carried forward)

• I8: No k ≤ ef validation
• I7: cudaFuncSetAttribute per-search caching
• M1 (unused #include <atomic>), M4 (ScratchSlot destructor missing cudaSetDevice)

Verified via Python simulation (no CUDA/GPU available for compile/test).

[premal]

Code Review — Delta 3 (k<=ef clamp, smem caching, distance convention)

Previously-Flagged Items — All Addressed

• [I8 FIXED] k <= ef: Auto-clamp ef = max(ef, k) + clear error if k can't fit in shared memory ✅
• [I7 FIXED] cudaFuncSetAttribute caching: Now uses static std::atomic<size_t> configured_smem per template instantiation ✅ (but see new issue below)
• [M4 FIXED] ScratchSlot destructor: cudaSetDevice(scratch.device) added before cudaStreamDestroy ✅
• [M1 FIXED] <atomic> include now actually used ✅

New Issue — Important (should-fix)

cudaFuncSetAttribute caching TOCTOU race: The static std::atomic with load-check-store pattern has a time-of-check-to-time-of-use race. Two concurrent searches with different ef values on the same template instantiation can interleave so the smaller-ef search's cudaFuncSetAttribute runs after the larger one, downgrading the global kernel attribute while configured_smem records the larger value. A future search with intermediate smem skips the set (thinks it's configured) but the actual attribute is too small → kernel launch failure.

Verified via Python simulation. Fix: CAS loop or mutex to ensure monotonic increase — only set the attribute if the new smem_size is strictly greater than the current configured_smem, and loop until CAS succeeds.

Distance Convention Change — Correct

The copy-out now negates IP/cosine back to true similarity (-result_dists[i]), L2 passed through. Sentinel padding is metric-aware (-FLT_MAX for IP/cosine empty, FLT_MAX for L2 empty). Upper-layer kernel correctly unchanged (uses internal negated convention, only produces entry points). Requires knowhere to stop negating on its side — confirmed done.

Degenerate Param Validation — Correct

if (ef <= 0 || sw <= 0) throws with a clear error message before any kernel launch. Good defensive check.

No new correctness bugs. Static analysis only (no CUDA/GPU for compile/test).

[premal]

Code Review — Delta 4 (TOCTOU mutex fix + entry_point validation + isCosine API)

All Previously-Flagged Should-Fix Items Now Resolved

• [TOCTOU race FIXED] cudaFuncSetAttribute caching now uses static std::mutex + std::lock_guard instead of the atomic load-check-store pattern. The mutex makes check-set-store atomic and monotonic — the kernel attribute only ever grows. Race fully eliminated. ✅

• [Entry-point validation NEW] Validates hnsw.entry_point < 0 || >= n_rows before GPU upload. Catches malformed CPU indexes with sentinel entry points (e.g. -1 → UINT32_MAX) that would cause OOB device reads. Error message includes the actual value and n_rows. ✅

• [isCosine()/useInnerProduct() API NEW] Clean public accessors that return the metric interpretation captured at copyFrom()/copyFromWithMetric() time. This is the faiss-side enabler for knowhere's I1 fix — knowhere can now derive the metric from the GPU index instead of search config. ✅

No New Issues

No critical or important issues. Minor: reset() doesn't clear is_cosine_/use_ip_ (stale values, but safe since searches guard on is_trained).

All previously-flagged should-fix items from our review series are now resolved. No blockers for merge.

[premal]

Code Review — Delta 5 (reset() clears metric fields)

✅ M1 from Delta 4 — FIXED

reset() now clears is_cosine_ and use_ip_ alongside deviceIndex_, ntotal, and is_trained. This addresses our delta 4 minor finding about stale metric fields surviving a reset. Clean fix.

No new issues. All previously-flagged items remain resolved.

[premal]

Delta 6 Review — faiss PR #1 (filtered search merged in)

The filtered search implementation from PR #7 has been folded into PR #1. Reviewed the delta against our previous PR #7 review.

Previously-Flagged Items

• I1 (cudaSetDevice in ensure_filter): ✅ FIXED — ensure_filter now starts with cudaSetDevice(device). Additionally, searchHost/searchHostInt8 set the device before calling, making it doubly-ensured.
• I2 (alpha gate order divergence): Not addressed — by design. Acknowledged in design doc §5.3, deferred to recall gate validation at 30-70% filter ratios.
• M1-M5 (minor): Not addressed — all non-blocking. Type punning (M1), unused meta[5] (M2), uninit invalid frontier (M3), no bounds check in is_bitset_filtered (M4), BF block size 128 (M5).
• C1 (TOCTOU race): ✅ Preserved — mutex fix correctly carried into per-<HAS_FILTER> instantiation.

New Issues (both minor)

• N1 (Minor): The unfiltered initial-seed path has one extra __syncthreads() (after the if constexpr block) that parallel_merge_into_result already provides. Functionally harmless (~few cycles), but means the "codegen-identical" comment is slightly inaccurate. The main loop does this correctly (sync is inside the if constexpr branch only) — inconsistency is only in the initial seed section.
• N2 (Trivial): searchHostInt8's dim%4 fallback to searchHost skips device setup, but searchHost does its own — safe.

Verdict

No blockers. Ready for compilation and the full recall/sanitizer gate. The I1 fix is the key item and it's resolved. The integration is clean — HAS_FILTER template, smem budget (36 B/ef filtered, 24 B/ef unfiltered), bitset upload, BF fallback launch orchestration, and destructor cleanup are all correct.

Full review at $M/reports/review-faiss-pr1-delta6.md

[devin-ai-integration]

Superseded: we re-cut clean gpu-hnsw branches from v2's base. Current PR is #11. Closing; branch retained.