Skip to content

kind/feat: GPU HNSW index with OCQ kernel and INT8 support#1686

Open
premal wants to merge 24 commits into
zilliztech:mainfrom
6si:gpu-hnsw-sq
Open

kind/feat: GPU HNSW index with OCQ kernel and INT8 support#1686
premal wants to merge 24 commits into
zilliztech:mainfrom
6si:gpu-hnsw-sq

Conversation

@premal

@premal premal commented Jun 19, 2026

Copy link
Copy Markdown

Summary

Adds GPU_HNSW index type for GPU-accelerated HNSW search on INT8 vectors, enabling Milvus deployments with NVIDIA GPUs to serve HNSW queries entirely on-device.

  • Implements Overflow Candidate Queue (OCQ) for correct HNSW graph traversal on GPU — without OCQ, beam-width search produces incorrect results at large ef values
  • Stores INT8 vectors directly on GPU (4x VRAM reduction vs FP32)
  • Eager GPU upload on Deserialize; eliminates per-search allocations and reduces sync overhead
  • Warp-cooperative distance computation for layer0 search
  • SMEM overflow clamped to avoid silent kernel failures at large ef + high-dim vectors
  • Fixes correct QT_8bit_direct_signed decode bias on GPU upload
  • Registers StaticCreateConfig and StaticHasRawData for full Knowhere lifecycle compatibility

Test plan

  • Multi-segment lifecycle (absorption) tests pass
  • INT8 bias decode and smem overflow clamp unit tests pass
  • 8/8 GPU HNSW integration tests pass
  • Benchmarked at 2M × 1024-dim INT8 vectors: 8,500 vec/s, 99.8% recall

🤖 Generated with Claude Code

@sre-ci-robot

Copy link
Copy Markdown
Collaborator

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: premal
To complete the pull request process, please assign zhengbuqian after the PR has been reviewed.
You can assign the PR to them by writing /assign @zhengbuqian in a comment when ready.

The full list of commands accepted by this bot can be found here.

Details Needs approval from an approver in each of these files:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

@sre-ci-robot

Copy link
Copy Markdown
Collaborator

Welcome @premal! It looks like this is your first PR to zilliztech/knowhere 🎉

@mergify

mergify Bot commented Jun 19, 2026

Copy link
Copy Markdown

@premal 🔍 Important: PR Classification Needed!

For efficient project management and a seamless review process, it's essential to classify your PR correctly. Here's how:

  1. If you're fixing a bug, label it as kind/bug.
  2. For small tweaks (less than 20 lines without altering any functionality), please use kind/improvement.
  3. Significant changes that don't modify existing functionalities should be tagged as kind/enhancement.
  4. Adjusting APIs or changing functionality? Go with kind/feature.

For any PR outside the kind/improvement category, ensure you link to the associated issue using the format: “issue: #”.

Thanks for your efforts and contribution to the community!.

premal and others added 21 commits June 19, 2026 03:29
…SW search

Implements GPU HNSW search in knowhere by converting CPU-built HNSW indexes
to run on GPU at search time. No changes to the CPU index build pipeline.

- GPU_HNSW_SQ: supports SQ8/FP16/BF16 storage (IndexScalarQuantizer)
- GPU_HNSW: supports plain F32 storage (IndexFlat)
- All three metrics supported: L2, IP, COSINE
- Kernels copied from cuvs into knowhere (no cuvs fork needed)
- Lazy GPU index build on first Search(), mutex-protected
- Bitset filter guard (returns error if non-empty — not yet supported)
- IP/COSINE distances negated back to positive after kernel returns

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…HNSW

Two root-cause bugs prevented GPU HNSW from running:

1. GetFaissHnswIndex() was dynamic_casting to ::faiss::IndexHNSW* but
   knowhere uses faiss::cppcontrib::knowhere::IndexHNSW which inherits
   from faiss::Index, not ::faiss::IndexHNSW. Cast always returned null.
   Fixed all interface and build files to use the correct type.

2. CMAKE_CUDA_ARCHITECTURES was hardcoded to 75-real (Turing) but the
   target GPU is L40S (Ada, sm_89). Changed to 89-real; all 15 tests pass.

Also adds proper cudaGetLastError() checks after kernel launches.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Single GpuHnswIndexNode handles all CPU storage formats (F32 and SQ8)
transparently. Key aliasing at Deserialize time accepts both "HNSW" and
"HNSW_SQ" CPU binaries. Removes GPU_HNSW_SQ constant and redundant
test file. All 13 tests pass.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Builds fat binary supporting Turing through Blackwell (CC 7.0–12.0),
required for g7e.2xlarge (Blackwell CC 12.0).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
- gpu_hnsw_faiss_build.hpp: log n_rows, entry_point, max_level, per-layer
  node counts, ep_in_layer flag, and neighbor occupancy (first 10 nodes)
  to help diagnose recall collapse hypotheses A/B
- gpu_hnsw_impl.cuh: after upper_layer_search_kernel, memcpy entry points
  back to host and log unique count vs global_ep to diagnose if upper-layer
  search is producing diverse entry points
- faiss_hnsw.cc: register GPU_HNSW for int8 input via IndexNodeDataMockWrapper
  (converts int8 queries to fp32 before delegating to GpuHnswIndexNode)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
- gpu_hnsw_impl.cuh: wrap entry point diversity readback in `static bool
  logged_ep_diversity` so the D2H memcpy + stream sync runs only on the
  first search call (~100μs one-time cost instead of per-call overhead)
- gpu_hnsw_impl.cuh: clamp ef to max that fits in 48KB shared memory
  (formula: (49152 - sw*max_degree0*8 - sw*4 - 12) / 12) before
  calc_layer0_smem_size; logs "[gpu_hnsw] clamping ef X->Y" if triggered.
  Prevents kernel crash at ef >= ~4088 with default sw=4, max_degree0=32.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Production fix:
- Clamp ef to prevent shared memory overflow crash at ef>=4096
  (smem formula ef*12+overhead exceeds 48KB GPU limit)

Diagnostics (compiled out unless GPU_HNSW_DIAGNOSTICS is defined):
- Build-time: graph structure, layer-0 neighbor sanity, upper layer occupancy
- Search-time: entry point diversity, beam search iteration tracking,
  brute-force sample distance check
- Interface: SQ type, metric type, storage type logging

Signed-off-by: Devin AI <devin@6sense.com>
Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…docstring

Add two absorption tests simulating Milvus segment lifecycle:
- F32 COSINE: two 5K segments -> merged 10K segment, verify recall
  and that merged recall >= per-segment recall
- SQ8 L2: same lifecycle + GPU cache reset on segment reload,
  verify deterministic results after re-deserialization

Also fix duplicate docstring comment on calc_layer0_smem_size.
clang-format applied to touched files.

Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
The sync was accidentally left outside the #ifdef block, causing an
unconditional GPU sync on every search call (adds ~0.5-1ms latency).
Move it inside the diagnostic block where it is only needed for the D2H
memcpy of entry points.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Optimizations to reduce GPU HNSW search latency:

1. Pre-allocated scratch buffers (search_scratch in gpu_hnsw_index):
   - d_queries, d_neighbors, d_distances, d_entry_points, d_visited_bitmaps
   - Allocated once on first search, reused across calls (grow-only)
   - Eliminates ~10 cudaMalloc/cudaFree calls per search (~1-2ms overhead)
   - Thread-safe via scratch_mutex (serializes concurrent searches)

2. Pre-built d_upper_layer_ptrs at index construction time:
   - Upper layer pointer array uploaded once during build_gpu_index
   - Eliminates per-search cudaMalloc + cudaMemcpyAsync + cudaFree for layer ptrs

3. Dedicated non-blocking CUDA stream (cudaStreamNonBlocking):
   - Created at build time, stored in gpu_hnsw_index::search_stream
   - Replaces null (default) stream which serializes with all GPU operations
   - Enables async memcpy without default-stream serialization

4. Removed unnecessary cudaStreamSynchronize between upper layer and L0:
   - Both kernels launch on same stream, so ordering is implicit
   - The sync was only needed for the diagnostic D2H memcpy (now inside #ifdef)
   - Saves ~1ms per search call in production builds

5. Async H2D transfer (cudaMemcpyAsync for query upload):
   - Query vectors copied to GPU asynchronously on the dedicated stream
   - CPU thread is free during transfer

6. CUDA event timing (gated behind GPU_HNSW_DIAGNOSTICS):
   - Records H2D, kernel, and total time for first 5 searches
   - Enable with -DGPU_HNSW_DIAGNOSTICS to profile the search pipeline

Before: ~21 synchronizing CUDA API calls per search
After:  1 cudaStreamSynchronize + 2 cudaMemcpy (D2H results)

Signed-off-by: Devin AI <devin@cognition.ai>
Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
For QT_8bit_direct_signed SQ indexes (HNSW_int8), upload raw INT8 codes
to GPU instead of dequantizing to float32. This reduces the dataset
memory footprint by 4x:

  2M vectors, dim=384:  3.07 GB (fp32) → 768 MB (int8)
  6M vectors, dim=384:  9.22 GB (fp32) → 2.30 GB (int8)

Implementation:
- Template search kernels (upper_layer + layer0_beam) on DataT
- Add load_elem(int8_t*) overload for on-the-fly int8→float conversion
- For COSINE metric: precompute per-vector reciprocal L2 norms at build
  time, multiply into IP distance in kernel. This avoids pre-normalizing
  vectors (which would require float storage). Norm array is n_rows × 4B
  (negligible vs 4x dataset savings).
- Detect QT_8bit_direct_signed at build time; fall back to float32
  dequantization for other SQ types (QT_8bit, QT_fp16, etc.)
- Refactor upload_to_gpu into upload_graph_to_gpu + upload_fp32_dataset
  / upload_int8_dataset for cleaner separation

Kernel distance computation with inv_norms:
  cosine_dist(q_norm, v) = -dot(q_norm, v/|v|)
                         = -dot(q_norm, v) * inv_norm[v]
                         = thread_ip_distance(q, v_int8) * inv_norm[v]

For L2 metric: no norms needed, direct int8→float per-element.

Signed-off-by: Devin AI <devin@cognition.ai>
Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
QT_8bit_direct_signed encodes values as (original + 128) stored in uint8.
The upload_int8_dataset function was copying raw biased bytes to GPU and
reading them through int8_t* in the kernel, which gives wrong values:
e.g., original -50 stored as byte 78, kernel reads int8_t 78 instead of -50.

Fix: convert biased-uint8 to true signed int8 on host before GPU upload.
signed_codes[i] = (int8_t)(codes[i] - 128) recovers the original value.
The kernel's load_elem(int8_t*) then correctly returns the decoded float.

Also fixes inv_norms computation which had the same bias error.

This restores correct distance computation for COSINE metric with INT8
vectors, fixing the 0% recall regression introduced in 1b9e346.

Co-Authored-By: Devin AI <noreply@cognition.ai>
Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Two new regression tests for GPU HNSW:

1. INT8 bias decode correctness (COSINE, 384-d, N=10K):
   Verifies QT_8bit_direct_signed codes (biased uint8, code = original + 128) are
   correctly converted to signed int8 before GPU upload (commit 60ca031).
   Before the fix: kernel reads biased values → garbage distances → R@1 = 0%.
   Test requires R@1 >= 0.90 vs brute-force and self-match >= 0.80.

2. ef overflow smem clamp:
   Verifies that ef values exceeding shared memory capacity (e.g. ef=8192 when
   max_ef ~4073 for sw=1, max_degree0=32) are clamped rather than causing CUDA
   launch failures. Requires valid results with recall >= 0.85.

Both tests require GPU to run. Tagged [int8_bias] and [smem_clamp] respectively.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…HNSW search

Adds a secondary expansion buffer (overflow queue) in global memory that
captures candidates rejected from the result buffer but still worth
expanding. This makes the GPU kernel a correct HNSW implementation
matching CPU HNSW semantics:

- Result buffer (Tier 1, shared memory): top-ef candidates, expanded first
- Overflow queue (Tier 2, global memory): candidates ranked ef+1..overflow_ef

Unified loop selects parents from result buffer first, falls through to
overflow when result buffer is exhausted. Overflow expansion can discover
candidates that re-enter the result buffer (phase oscillation).

Key changes:
- search_params: add overflow_factor (default 3, overflow_ef = 3*ef)
- search_scratch: add overflow queue device buffers (ids, dists, expanded, count)
- layer0_beam_search_kernel: unified parent selection + overflow spill on eviction
- overflow_insert: sorted insertion helper for the overflow queue

Memory overhead: ~4.5 KB/query at ef=128 (vs 1MB/query for visited bitmap).
Convergence: bounded by (ef + overflow_ef) / search_width + 20 iterations.

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…-search

Move build_gpu_index() call from Search() into Deserialize() so the GPU
index is ready immediately when a segment loads. Eliminates cold-start
latency on first query after segment load/restart.

The lazy init in Search() remains as a fallback for edge cases where
Deserialize was called without metric config.

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…ly loads

Two throughput optimizations:

1. Reduce overflow_factor from 3 to 2: fewer overflow expansions and
   iterations (max_iter drops from (ef+3*ef)/sw+20 to (ef+2*ef)/sw+20).
   At 99.8% recall there is headroom to trade minor recall for throughput.

2. Add __ldg() intrinsic to all read-only global memory loads:
   - Dataset vector elements (load_elem helpers)
   - Graph neighbor IDs (d_layer0_graph)
   - Inverse norms (d_inv_norms)
   - Upper layer node IDs and neighbors
   __ldg routes through the texture/read-only cache, which handles
   random access patterns better than L1/L2 for read-only data at scale.

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Restructure neighbor expansion to use cooperative distances: multiple
threads collaborate on a single distance computation, each handling a
slice of dimensions and reducing via warp shuffle.

threads_per_dist is selected at runtime based on dimension:
  dim >= 256: 4 threads per distance (96+ dims/thread for dim=384)
  dim >= 96:  2 threads per distance
  dim < 96:   1 thread (no cooperation, original behavior)

Key benefits at large N (32.9M+ vectors):
- Fewer unique cache lines accessed per warp: 128/4 = 32 vectors
  simultaneously (vs 128 with 1-thread-per-dist), reducing L2 thrashing
- Sub-warp group masks allow groups to diverge independently (no
  full-warp synchronization required for bitmap/validity checks)
- Group lane 0 handles graph reads and bitmap_visit atomics; all lanes
  participate in distance computation; only lane 0 writes staging

Applied to both seeding phase (entry point neighbors) and main loop
Step 2 (parent expansion). Upper-layer search unchanged (already uses
full-warp parallelism across neighbors with 1 lane per distance).

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Three compute-bound optimizations targeting 20K QPS at scale:

1. Parallel bitonic sort of staging buffer (all 128 threads)
   Before: thread-0 serial O(staging_count × ef) merge per iteration
   After: parallel sort O(log²n) then thread-0 early-exit merge
   Sorted staging allows early break once entries stop beating the
   result buffer, reducing thread-0 serial work by ~80%.

2. Early convergence (stale iteration detection)
   Track consecutive iterations where result_dists[ef-1] doesn't
   improve. Break after 4 stale iterations even if overflow has
   unexpanded entries. Eliminates wasted iterations where the search
   has converged but the main loop continues expanding unpromising
   overflow candidates.

3. search_width 4→8
   More parents per iteration → fewer total iterations to converge.
   With 8 parents × 32 max_degree = 256 neighbors expanded per
   iteration (vs 128 before), the graph is explored ~2x faster.
   max_iterations auto-adjusts: (64+128)/8+20 = 44 (was 68).

Combined with early convergence, typical iteration count drops from
~50-60 to ~15-25 while exploring the same effective neighborhood.

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…nvergence

The bitonic sort + search_width=8 regressed throughput by 8-18% at
moderate-high concurrency. Root cause: 36 extra __syncthreads barriers
per iteration from the sort outweigh the early-exit benefit in the
merge loop. search_width=8 doubled staging_count (256 vs 128), making
each iteration heavier despite fewer total iterations.

Revert to search_width=4 and serial merge (v41 baseline). Keep only
the early convergence heuristic: meta[3] tracks consecutive iterations
where result_dists[ef-1] doesn't improve; break after 4 stale
iterations. This is zero-overhead (single int comparison) and cuts
wasted post-convergence iterations without adding barriers.

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…arly convergence

Documents the full GPU HNSW kernel architecture:
- Overflow Candidate Queue (OCQ) for correct HNSW semantics
- Warp-cooperative distance computation (sub-warp groups)
- Early convergence heuristic (stale iteration detection)
- Phase 1/Phase 2 search pipeline
- Memory layout (shared + global)
- Configuration parameters and tuning guidelines
- Performance history (v38-v43) with benchmarks

Also adds GPU_HNSW to the index-types skill reference.

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
…ate config' warning

GPU_HNSW was registered via KNOWHERE_REGISTER_GLOBAL (IndexFactory only)
but not KNOWHERE_REGISTER_STATIC (IndexStaticFaced). This caused
CreateConfig() in index_static.cc:43 to fall through to the warning
path on every segment load.

Add explicit RegisterStaticFunc<GpuHnswIndexNode> calls for both fp32
and int8 data types, which populates staticCreateConfigMap with
GpuHnswIndexNode::StaticCreateConfig (returns FaissHnswConfig).

Signed-off-by: Devin AI <devin@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
premal and others added 2 commits June 19, 2026 03:29
- Bitset fix: check !bitset.empty() && bitset.count() > 0 to avoid
  false filtered-search rejection on fresh collections with all-zero
  bitsets (causes GPU_HNSW to refuse valid unfiltered queries)
- Static constructor: replace inline const global registrations with
  __attribute__((constructor)) to fix static init order issues that
  caused "unhandled create config for indexType: GPU_HNSW" warnings

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Premal Shah <premal@6sense.com>
Add StaticHasRawData returning true to GpuHnswIndexNode so the segment
loader knows the index contains raw vectors and skips loading redundant
field binlogs into CPU RAM.  Without this, each querynode loads raw
vectors both inside the CPU HNSW structure (retained for GetVectorByIds)
and as separate field data — roughly doubling CPU memory usage.

Also register GPU_HNSW in IndexStaticFaced<int8> so that int8 segments
no longer trigger the 'unhandled create config for indexType: GPU_HNSW'
warning on the load path (CGO_DYN context).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Premal Shah <premal@6sense.com>
@mergify mergify Bot added dco-passed and removed needs-dco labels Jun 19, 2026
@premal premal changed the title feat: GPU HNSW index with OCQ kernel and INT8 support kind/feat: GPU HNSW index with OCQ kernel and INT8 support Jun 19, 2026
Lead with feature explanation (what, why, key properties, performance)
before design details. Add graph conversion section, dependencies
link, and clearer section organization.

Co-Authored-By: Devin AI <noreply@devin.ai>
Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@mergify

mergify Bot commented Jun 19, 2026

Copy link
Copy Markdown

@premal e2e jenkins job failed, comment /run-e2e-sse can trigger the job again.

@mergify

mergify Bot commented Jun 19, 2026

Copy link
Copy Markdown

@premal e2e jenkins job failed, comment /run-e2e can trigger the job again.

@foxspy

foxspy commented Jun 26, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the contribution and for the detailed work on this PR.

After internal discussion, we do not plan to accept PRs that add or maintain GPU index implementations directly in Knowhere in this form. Going forward, GPU-related capabilities should be integrated through upstream FAISS/cuVS whenever possible, rather than carrying a separate implementation in this repository.

If this work is generally useful, we would suggest proposing the corresponding changes to FAISS or cuVS upstream. Once the capability is available there, Knowhere can pick it up through the normal dependency upgrade path.

Given that direction, we will not move forward with this PR. Thanks again for the effort and for sharing the implementation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants