Skip to content

[DRAFT][Feature] Add scalar quantization support for vector extension#28

Open
rahul-iyer wants to merge 1 commit into
LadybugDB:mainfrom
rahul-iyer:quantized-rerank-cache
Open

[DRAFT][Feature] Add scalar quantization support for vector extension#28
rahul-iyer wants to merge 1 commit into
LadybugDB:mainfrom
rahul-iyer:quantized-rerank-cache

Conversation

@rahul-iyer

Copy link
Copy Markdown
Contributor

Adds scalar quantization support for HNSW vector indexes behind the quantization option.

This PR supports:

  • quantization := 'sq8'
  • quantization := 'sq16'
  • quantized in-memory distance computation for ANN traversal
  • optional full-precision reranking via use_full_precision_rerank
  • SIMD-friendly SQ16 dot-product kernel with scalar fallback
  • quantized query-vector/cache alignment fixes
  • visibility checks for quantized on-disk embeddings
  • quantized cache invalidation on index mutation

Dot product with quantization is rejected for now because the quantized metric path currently supports
L2/L2SQ/Cosine-style scoring.

Behavior

Without quantization, HNSW continues to use the existing full-precision float embedding path.

With quantization enabled:

  • index traversal uses quantized embeddings
  • query vectors are quantized once per query
  • cached quantized payloads store compact integer vectors plus scale/norm metadata
  • if use_full_precision_rerank := true, final candidates are rescored using full-precision embeddings
  • if use_full_precision_rerank := false, results are returned directly from the quantized metric path

Benchmark

Dataset: real-world OpenAI/DBPedia-style benchmark subset
Metric: recall@10 against full-precision ground truth
Query workload: repeated ANN search over loaded benchmark DB

With Cache Embeddings = true
base load=7.145s index=72.196s ingest=79.342s ms/query=19.812 db_mib=606.4 peak_rss_mib(load/index/query/acc)=9140.5/1366.2/751.1/725.3 recall@10=0.9960

sq8 load=6.678s index=10.762s ingest=17.439s ms/query=15.544 db_mib=606.5 peak_rss_mib(load/index/query/acc)=9138.5/812.1/823.0/803.0 recall@10=0.9460

sq16 load=4.338s index=13.876s ingest=18.215s ms/query=16.214 db_mib=606.4 peak_rss_mib(load/index/query/acc)=9599.7/883.7/896.1/876.2 recall@10=0.9950

Observations

SQ8 gives the largest indexing speedup and lower query latency, with an expected recall tradeoff.

SQ16 keeps recall very close to full precision while still reducing index build time substantially.

The biggest win is index construction time:

  • SQ8 index build is roughly 11x faster than full precision in this benchmark.
  • SQ16 index build is roughly 5x faster than full precision.
  • Query latency improves modestly for both SQ8 and SQ16.

Correctness Notes

This PR also fixes several issues found during testing:

  • quantized query buffer sizing/alignment
  • SQ8/SQ16 cached metadata layout
  • traversal using the quantized metric function instead of the full-precision metric
  • stale quantized cache after index mutation
  • visibility checks for quantized on-disk embeddings

Quantized cache invalidation is currently coarse-grained: mutations invalidate the full quantized cache.
A future improvement should make this transaction-aware and row/version based.

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