Skip to content

fix(gpu_hnsw): bound visited-bitmap VRAM via nq-chunked layer-0 search#9

Closed
devin-ai-integration[bot] wants to merge 1 commit into
gpu-hnsw-faissfrom
devin/gpu-hnsw-oom-chunk
Closed

fix(gpu_hnsw): bound visited-bitmap VRAM via nq-chunked layer-0 search#9
devin-ai-integration[bot] wants to merge 1 commit into
gpu-hnsw-faissfrom
devin/gpu-hnsw-oom-chunk

Conversation

@devin-ai-integration

@devin-ai-integration devin-ai-integration Bot commented Jul 20, 2026

Copy link
Copy Markdown

Summary

Fixes a device-memory OOM in GPU HNSW search caused by the grow-only search scratch. The per-slot visited bitmap is nq * ceil(N/32) * 4 bytes and, once grown, is never released. Under high search concurrency (one bitmap per concurrent pool slot) on large segments this exhausts VRAM — 16 concurrent batch=512 searches on a ~538M-row segment grew the pool to ~97 of ~98 GB and OOM'd every subsequent allocation (surfaced during the v90 Blackwell validation).

Fix: process queries in nq-chunks so the bitmap is bounded regardless of batch size / concurrency. No kernel changes — every per-query buffer (queries, entry points, neighbors, distances, visited bitmap, BF worklist) is already indexed by the chunk-local block index (blockIdx.x / worklist entry), so the host just launches per chunk against caller-offset base pointers. Queries are independent in HNSW, so chunking is result-invariant.

chunk_nq = gpu_hnsw_bitmap_chunk(nq, N)   // = nq when the batch fits the cap
for c0 in 0, chunk_nq, 2*chunk_nq, ...:
    cnq  = min(chunk_nq, nq - c0)
    q0   = d_layer0_queries + c0*dim       // int8 on DP4A path, else fp32
    fq0  = d_queries        + c0*dim       // fp32, upper-layer descent
    ep0  = d_entry_points   + c0
    nb0  = d_neighbors      + c0*k
    ds0  = d_distances      + c0*k
    <upper-layer greedy descent>(fq0, ep0, cnq)
    memset(visited_bitmap, 0, calc_visited_bitmap_size(cnq, N))
    layer0_beam_search<<<cnq, ...>>>(q0, ep0, nb0, ds0, cnq, ...)   // + per-chunk BF fallback
  • gpu_hnsw_bitmap_chunk(nq, N) derives the per-launch query count from a VRAM cap (default 256 MiB, env-tunable via GPU_HNSW_BITMAP_MB, or GPU_HNSW_BITMAP_BYTES for byte-precise control in tests). Always in [1, nq]; returns nq when the batch already fits, so small segments run a single pass — behavior identical to before.
  • GpuHnswSearchScratch::ensure() now sizes the bitmap for one chunk (bm_nq) instead of the full batch — this is the actual allocation that no longer grows unbounded. The cap only ever lowers the chunk vs. what ensure() sized, so the buffer is always large enough for the launched chunk.
  • The shared-memory sizing / ef finalization block is query-count independent and is hoisted to run once before the chunk loop; the cudaFuncSetAttribute high-water tracking (per <DataT,QueryT,USE_DP4A,HAS_FILTER>) is unchanged.

All existing semantics preserved: FP32 / FP16 / BF16 / generic-INT8 / native-INT8-DP4A, L2 / IP / COSINE, filtered search (two-tier beam + alpha gate + per-query BF fallback), up-front brute force, valid IDs / sentinels, CPU-HNSW parity.

Parity coverage lands with the knowhere re-vendor (this faiss tree only builds inside knowhere): a [gpu_hnsw_chunk_parity] test runs each search twice (default cap = single pass vs. GPU_HNSW_BITMAP_BYTES forced tiny = many chunks) and asserts identical ids/distances across fp32 & int8-DP4A, unfiltered / mid-ratio filtered / up-front-BF (99%) for L2/IP/COSINE.

Related PRs

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

The layer-0 search scratch is grow-only per pool slot: the visited bitmap is nq * ceil(N/32) * 4 bytes and, once grown, is never released. Under high search concurrency (one bitmap per concurrent pool slot) on large segments this exhausted device memory -- 16 concurrent batch=512 searches on a ~538M-row segment grew the pool to ~97 of ~98 GB and OOM'd every subsequent allocation.

Process queries in nq-chunks so the bitmap is bounded regardless of batch size / concurrency. gpu_hnsw_bitmap_chunk() derives the per-launch query count from a VRAM cap (default 256 MiB, env-tunable via GPU_HNSW_BITMAP_MB / GPU_HNSW_BITMAP_BYTES); GpuHnswSearchScratch::ensure() sizes the bitmap for one chunk and the search launches per chunk with caller-offset base pointers. Queries are independent in HNSW search and every per-query buffer is indexed by the chunk-local block index, so chunking is result-invariant (no kernel changes). Small segments yield chunk == nq -> a single pass, identical to prior behavior.

Signed-off-by: premal <premal@6sense.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@premal premal self-assigned this Jul 20, 2026
@devin-ai-integration

Copy link
Copy Markdown
Author

🤖 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

Copy link
Copy Markdown
Author

Superseded by the current gpu-hnsw stack (#11). The nq-chunked visited-bitmap OOM fix is included there. Closing; branch retained.

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