diff --git a/BENCHMARK.md b/BENCHMARK.md index a8b5fb3c..c54a0cd7 100644 --- a/BENCHMARK.md +++ b/BENCHMARK.md @@ -878,6 +878,135 @@ indexing); Moon still ingests 5.8–10.4× faster. (2) recall@10 dips 0.9992 → the index now serves from 3 segments instead of one (multi-segment beam truncation), a ~0.001 recall cost for the 2.3–4.6× faster time-to-green; still within 0.002 of Qdrant. +### 10.10 2026-07-08 ANN-benchmarks 4-way: Moon vs RediSearch vs Qdrant vs turbovec (GCE ARM, 1.18M/1M vectors) + +First campaign on the **standard ANN-benchmarks datasets with bundled ground truth** +(previous sections used synthetic/MiniLM corpora ≤50K). Full recall/QPS trade-off curves +via the new runtime `FT.CONFIG SET EF_RUNTIME` (commit `0afac2c4`) — one ingest, +one settled index, N query points. + +**Environment:** GCE t2a-standard-8 (ARM Ampere, 8 vCPU, us-central1-a), Moon +`0afac2c4` (feat/vector-fastscan, `-C target-cpu=native`, monoio, `--shards 8 +--appendonly no --disk-offload disable --max-unflushed-immutable-segments 0`); +RediSearch = `redis/redis-stack-server:latest` (Docker, HNSW M=16 EF_CONSTRUCTION=200); +Qdrant = `qdrant/qdrant:latest` (Docker, default HNSW, wait-for-optimizer green); +turbovec = github.com/RyanCodrai/turbovec (in-process Rust/Python TurboQuant flat index). +Harness: redis-py / qdrant-client / native API, 1000 test queries, recall@10 vs bundled +ground truth, qps1/qps8 = 1/8 client threads (no pipelining). Moon "settled" = ingest → +green → `VACUUM VECTOR` merge loop to 1 segment/shard (8 total). + +#### glove-200-angular — 1,183,514 × 200d, COSINE, k=10 + +| system | config | recall@10 | qps (1 client) | qps (8 clients) | +|--------|--------|:---------:|:----:|:----:| +| **Moon** (SQ8, settled) | ef=16 | 0.648 | **1448** | **2907** | +| | ef=64 | 0.827 | **888** | **1899** | +| | ef=128 | 0.887 | **612** | **1309** | +| | ef=256 | 0.933 | **400** | **876** | +| RediSearch | ef=16 | 0.496 | 2328 | 3618 | +| | ef=64 | 0.694 | 1202 | 1704 | +| | ef=128 | 0.771 | 816 | 1074 | +| | ef=256 | 0.836 | 524 | 602 | +| Qdrant | ef=16 | 0.570 | 410 | 574 | +| | ef=64 | 0.768 | 335 | 544 | +| | ef=256 | 0.887 | 214 | 426 | +| turbovec | 4-bit | 0.875 | 63 | 62 | +| | 2-bit | 0.625 | 105 | — | + +**Iso-recall verdicts (the only fair read):** + +- **vs RediSearch @ ~0.83 recall:** Moon 888/1899 vs 524/602 — **1.7× / 3.2× win**. + (RediSearch's 2328-qps ef=16 row is at 0.50 recall — not a comparable operating point.) +- **vs Qdrant @ ~0.887 recall:** Moon 612/1309 vs 214/426 — **2.9× / 3.1× win**. +- **vs turbovec @ ~0.875 recall:** Moon ~700+/1500+ vs 63/62 — **>10× win**; the flat + scan is O(N) and collapses at 1.18M vectors. + +**Time-to-ready** (same dataset): Moon ingest 165 s → **green (HNSW-tier, 0.85+ recall +available) at 277 s**; RediSearch 1498 s; Qdrant ~390 s. Moon's *fully-settled* state +(merge to 1 segment/shard) took a further ~61 min — see trade-offs below. + +#### glove subset 100K (turbovec's native scale), COSINE + +k=64 (turbovec's README config): + +| system | config | recall@64 | qps1 | qps8 | +|--------|--------|:---------:|:----:|:----:| +| **Moon** (SQ8, settled) | ef=64 | 0.886 | 647 | **1318** | +| | ef=128 | 0.952 | 524 | 1103 | +| | ef=512 | 0.998 | 272 | 639 | +| turbovec | 4-bit | 0.895 | **760** | 781 | +| | 2-bit | 0.663 | 1191 | 1243 | + +k=10: + +| system | config | recall@10 | qps1 | qps8 | +|--------|--------|:---------:|:----:|:----:| +| **Moon** (SQ8, settled) | ef=16 | 0.765 | 1787 | 3211 | +| | ef=64 | **0.930** | **1076** | **2286** | +| | ef=256 | 0.995 | 527 | 1139 | +| turbovec | 4-bit | 0.877 | 787 | 802 | + +At k=10 Moon wins both axes (higher recall AND higher qps1, 2.9× qps8). At k=64 the +in-process flat scan edges Moon on single-client qps (no network hop, cost amortized +over large k) while Moon wins 8-client throughput 1.7× and is the only one that can +reach >0.9 recall. turbovec gets no concurrency scaling (760 → 781). + +#### gist-960-euclidean — 1,000,000 × 960d, L2, k=10 + +| system | config | recall@10 | qps1 | qps8 | +|--------|--------|:---------:|:----:|:----:| +| **Moon** (SQ8, settled) | ef=16 | 0.816 | **995** | **2077** | +| | ef=64 | 0.962 | **528** | **1085** | +| | ef=256 | **0.994** | **216** | **472** | +| Qdrant | ef=16 | 0.624 | 264 | 390 | +| | ef=64 | 0.861 | 215 | 371 | +| | ef=256 | 0.965 | 134 | 314 | +| RediSearch | ef=16 | 0.486 | 1242 | 1827 | +| | ef=64 | 0.757 | 591 | 798 | +| | ef=256 | 0.927 | 246 | 270 | + +(turbovec omitted: IP/COSINE-only library, and its O(N) flat scan is non-viable at 1M.) + +**Iso-recall verdicts — Moon wins every band:** + +- **@ ~0.86:** Moon (ef≈24, interp.) ~**4× / ~5×** Qdrant's 215/371. +- **@ ~0.96:** Moon 528/1085 vs Qdrant 134/314 — **3.9× / 3.5×**. +- **@ ~0.93:** Moon *at higher recall* (0.962 @ 528/1085) vs RediSearch 0.927 @ + 246/270 — **2.1× / 4.0×**. +- Moon is the only system measured past 0.99 recall (0.994 @ 216/472 — still faster + than both competitors' ~0.93–0.965 operating points). + +**Time-to-ready** (gist): Moon ingest 396 s → green 437 s; RediSearch 2871 s; +Qdrant ~620 s. Moon settle (merge to 1 segment/shard) took a further ~97 min — the +merge-parallelism gap noted below. + +⚠ **Moon TQ4 on gist recorded recall 0.002–0.003 — a bug, not noise.** TQ's ADC ranked +L2 by `sphere_dist·‖a‖²`, which is only rank-valid on the unit sphere; gist is +unnormalized L2 and the estimator collapses. **Fixed in this branch** (`14a07025`): all +TQ scoring paths reconstruct `‖a−q‖² = (‖a‖−‖q‖)² + ‖a‖‖q‖·d̂²`, and `FT.CREATE` L2 +indexes now **default to SQ8** (explicit TQ*+L2 warns). GCE re-verification with +explicit TQ4 post-fix (same instance/harness, unsettled index so recall-only): +ef=16 **0.784**, ef=64 **0.942**, ef=256 **0.986** — a ~350× recovery, within ~0.03 of +SQ8 (4-bit direction precision at 960d). The main table above is the SQ8 run. + +#### Findings that changed Moon along the way + +1. **`EF_RUNTIME` was FT.CREATE-frozen** → recall/QPS curves needed one full index + rebuild per point. Now runtime-tunable via `FT.CONFIG SET` (also fixed FT.CONFIG SET + being silently local-shard-only under monoio at shards>1). +2. **MA1 write-stall guard fires on TOTAL immutable segments (>20), not unflushed + ones** — a 1M+ bulk load sits above the threshold while merges lag, throttling + ingest to ~190 vec/s (24× slowdown) on idle hardware. Workaround + `--max-unflushed-immutable-segments 0` during bulk loads; semantics fix pending. +3. **Unmerged segments multiply query cost**: per-query work ≈ shards × segments × ef. + Pre-settle (57 segments) the same index ran 4–5× slower at the same ef. `VACUUM + VECTOR` merges to 1 segment/shard but is local-shard-only over the wire and took + ~56 min for 1.18M × 200d (and ~97 min for 1M × 960d) at only ~1.2 cores — + GraphUnion merge parallelism is the next optimization target. +4. **Multi-segment search unions independent beams** → higher recall ceiling at equal + ef (unsettled ef=256: 0.9865 vs settled 0.933). The settled index needs ef 512+ to + reclaim the >0.95 band — at far higher qps than the unsettled equivalent. + --- ## 11. Graph Engine diff --git a/CHANGELOG.md b/CHANGELOG.md index 3dcfb36e..cfceb01e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,108 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added — CLI/moon.conf defaults for vector + graph tuning knobs (PR #TBD) + +- **`--vector-ef-runtime` / `--vector-rerank-mult` / `--vector-exact-beam`** + set the server-wide starting values every NEW vector index is created with + (FT.CREATE), so recall-sensitive fleets configure once instead of issuing + FT.CONFIG per index. Per-index `FT.CONFIG SET` always overrides. Same + ranges as FT.CONFIG (ef 10-4096 or 0=auto, mult 1-64), validated loudly at + startup. All three work as `moon.conf` keys (`vector-exact-beam yes|no`). +- **`--graph-result-cache-entries` / `--graph-result-cache-bytes`** size the + per-graph Cypher result cache (previously hardcoded 256 entries / 4 MiB). +- Startup wiring is first-write-wins process state, installed by both the + binary entry and `run_embedded`; unit tests and library embedders that + never install defaults keep the exact pre-flag behavior. + +### Added — Recall knobs: FT.CONFIG RERANK_MULT + EXACT_BEAM (PR #TBD) + +- **`FT.CONFIG SET RERANK_MULT `** (1-64, default 4) deepens the HQ-1 + exact-rerank stage: the top `n·k` beam candidates are re-scored with true f16 + sidecar distances before top-k truncation, recovering true neighbors the + quantized ADC ranking dropped below the default 4·k cut. Cost ~`n·k·dim` f16 + decodes per segment. +- **`FT.CONFIG SET EXACT_BEAM ON|OFF`** (default OFF) navigates the HNSW + beam itself with exact f16 distances instead of quantized ADC estimates — + beam candidate *selection* becomes exact, so recall is graph-limited + (Qdrant-parity at equal ef) rather than quantization-limited. QPS cost grows + with dimension. Segments without an exact-rerank sidecar (pre-HQ-1 disk + reloads) silently keep the quantized beam. +- Both knobs are per-index, apply to the next FT.SEARCH (all search paths: + sync, yielding, worker-pool fan-out, FT.RECOMMEND, hybrid), broadcast to all + shards, and persist across restarts via a new v4 index-meta sidecar format + (v1-v3 sidecars load with defaults). + +### Added — Runtime-tunable EF_RUNTIME via FT.CONFIG (PR #TBD) + +- **`FT.CONFIG SET EF_RUNTIME `** adjusts the HNSW search beam width at + runtime without rebuilding the index (RediSearch parity). Range matches + FT.CREATE (10-4096); `0` restores the auto heuristic. Applies to the next + FT.SEARCH immediately and persists via the index meta sidecar. `FT.CONFIG GET + EF_RUNTIME` added alongside. +- **Fixed — FT.CONFIG SET was local-shard only under monoio**: at shards>1 the + setting silently applied to 1/N index partitions (AUTOCOMPACT, + COMPACTION_WEIGHT, MERGE_RECALL_TOLERANCE were equally affected). SET now + broadcasts to all shards like FT.CREATE; GET stays local. + +### Fixed — TQ ADC estimator is now metric-faithful on raw L2 (PR #TBD) + +- **TQ quantization ranked L2 queries by `sphere_dist·‖a‖²`** — the unit-sphere + distance between normalized directions scaled by the document norm. On + unnormalized L2 data this makes every small-norm vector look near regardless + of direction (gist-960: recall 0.002). All TQ scoring paths (HNSW beam + + budgeted variant, mutable brute-force + FastScan pre-filter bound, flat + scan, multibit brute-force, A2 decoded-L2) now reconstruct the true metric: + `‖a−q‖² = (‖a‖−‖q‖)² + ‖a‖·‖q‖·d̂²_sphere`. COSINE/IP scoring is unchanged. + Unit tests pin recall on varying-norm data at both the mutable-scan and + HNSW-beam levels (7+/10 and 6+/10 vs exact f32 ground truth; 4/10 and worse + before the fix). GCE re-verification on gist-960 with explicit TQ4: + recall@10 0.784/0.942/0.986 at ef 16/64/256 vs 0.002-0.003 pre-fix (~350x), + within ~0.03 of SQ8. + +### Changed — FT.CREATE L2 indexes default to SQ8 quantization (PR #TBD) + +- **`FT.CREATE ... DISTANCE_METRIC L2` without an explicit `QUANTIZATION` now + defaults to SQ8** instead of TQ4. TQ's norm-scaled ADC estimator assumes + unit-sphere metrics (COSINE/IP) and collapses on unnormalized L2 data — + recall 0.002 measured on gist-960-euclidean (1M × 960d). SQ8 is + metric-faithful on raw L2. COSINE/IP defaults are unchanged (TQ4). An + explicit `QUANTIZATION TQ*` + L2 is still honored but logs a warning. +- Benchmarks: `BENCHMARK.md` §10.10 — ANN-benchmarks 4-way campaign (Moon vs + RediSearch vs Qdrant vs turbovec on glove-200-angular 1.18M / gist-960 1M / + glove-100K). Tuning guide gains bulk-load (`--max-unflushed-immutable-segments + 0`), settle (`VACUUM VECTOR`), and runtime-`EF_RUNTIME` recipes. + +### Added — FastScan SIMD vector scan: NEON TBL kernel + live-path integration (PR #TBD) + +- **NEON TBL FastScan kernel** (`src/vector/distance/fastscan.rs`): register-resident + 4-bit LUT accumulation via `vqtbl1q_u8` for aarch64 — the primary platform previously + fell back to the scalar kernel. 32 candidate distances per block at ~2 ns/candidate + (128d), 17-21× faster than the per-candidate scalar ADC path. +- **AVX2 kernel overflow fix**: nibble-pair distances were added as u8 before widening, + silently wrapping for LUT pairs summing past 255 and corrupting rankings. All kernels + now widen to u16 before adding and accumulate with saturating adds (scalar/NEON/AVX2 + bit-identical across the full u8 LUT range; parity tests no longer mask LUTs to 0x7F). +- **Adaptive quantized LUT builder** (`build_quantized_lut`): FAISS-style per-coordinate + bias + global scale chosen so entries fit u8 AND the worst-case accumulated sum fits + u16 (no kernel saturation for in-range data), with f32 reconstruction (`acc/scale + + bias`). Replaces the legacy `precompute_lut`, which hardcoded the v1 1/sqrt(768) + `CENTROIDS` table (encode/search codebook asymmetry — the same recall bug class fixed + by codebook v2) and a fixed `LUT_SCALE` that could overflow both u8 entries and the + u16 accumulator. +- **Mutable-segment FastScan pre-filter** (`src/vector/segment/mutable.rs`): the MVCC + brute-force scan now maintains a FAISS-interleaved shadow of the TQ4 codes (32-vector + blocks, +padded_dim/2 B/vector) and screens candidates with the SIMD kernel; only + candidates whose sound lower bound (`approx − 0.5·padded_dim/scale`) beats the current + heap worst get the exact f32 ADC rescore. Results are bit-identical to the plain scan + (the bound is a true lower bound; saturation only under-estimates). Measured: top-10 + over 20K×128d 1.25 ms → 96 µs (~13×), 5K×768d 1.66 ms → 423 µs (3.9×) on Apple + Silicon. Engages above 64 entries (per-query LUT build costs ~2-8 µs); SQ8/A2 + collections and TQ-prod scoring keep the existing paths. +- New criterion bench `benches/fastscan_bench.rs` (block kernels, LUT build, end-to-end + mutable scan A/B); shadow-layout + FastScan-vs-plain equality tests (full scan, + chunked scan, MVCC visibility + bitmap-filter paths). + ## [0.6.0] — 2026-07-08 **Release highlights** (full detail in the sections below; "PR #TBD" entries diff --git a/Cargo.toml b/Cargo.toml index e1e88596..f68a9639 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -228,6 +228,10 @@ harness = false name = "sq8_adc_bench" harness = false +[[bench]] +name = "fastscan_bench" +harness = false + [[bench]] name = "hnsw_bench" harness = false diff --git a/benches/fastscan_bench.rs b/benches/fastscan_bench.rs new file mode 100644 index 00000000..38db3305 --- /dev/null +++ b/benches/fastscan_bench.rs @@ -0,0 +1,220 @@ +//! Criterion benchmarks for the FastScan block kernels. +//! +//! Compares three ways of computing approximate TQ4 ADC distances for a +//! 32-vector batch: +//! - `block_scalar`: the scalar FastScan reference (interleaved layout, u8 LUT) +//! - `block_dispatch`: the SIMD FastScan kernel selected for this machine +//! (NEON TBL on aarch64, AVX2 VPSHUFB on x86_64) +//! - `tq_adc_x32`: the current per-candidate scalar path (`tq_l2_adc_scaled` +//! with an f32 LUT-free form) called 32 times — what HNSW beam search and +//! the mutable-segment brute-force scan do today. +//! +//! Also measures `build_quantized_lut` (the per-query fixed cost). + +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use moon::vector::distance::fastscan; +use moon::vector::segment::ivf::BLOCK_SIZE; +use moon::vector::turbo_quant::codebook::scaled_centroids; +use moon::vector::turbo_quant::tq_adc::tq_l2_adc_scaled; +use std::hint::black_box; + +/// Padded dimensions to benchmark (128 stays 128, 384 pads to 512, 768 to 1024). +const PADDED_DIMS: &[usize] = &[128, 512, 1024]; + +fn make_query(padded_dim: usize, seed: u32) -> Vec { + let sigma = 1.0 / (padded_dim as f32).sqrt(); + let mut s = seed; + let mut q = Vec::with_capacity(padded_dim); + for _ in 0..padded_dim { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + q.push(((s as f32) / (u32::MAX as f32) * 4.0 - 2.0) * sigma); + } + q +} + +/// Random nibble-packed codes for 32 vectors, both vector-major (for the +/// per-candidate path) and interleaved (for FastScan). +fn make_codes(dim_half: usize, seed: u32) -> (Vec>, Vec) { + let mut s = seed; + let mut per_vec = vec![vec![0u8; dim_half]; BLOCK_SIZE]; + for cv in per_vec.iter_mut() { + for b in cv.iter_mut() { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + *b = (s >> 24) as u8; + } + } + let mut interleaved = vec![0u8; dim_half * BLOCK_SIZE]; + for (v, cv) in per_vec.iter().enumerate() { + for (d, &byte) in cv.iter().enumerate() { + interleaved[d * BLOCK_SIZE + v] = byte; + } + } + (per_vec, interleaved) +} + +fn bench_fastscan_block(c: &mut Criterion) { + fastscan::init_fastscan(); + let mut group = c.benchmark_group("fastscan_block_32vec"); + + for &padded_dim in PADDED_DIMS { + let dim_half = padded_dim / 2; + let centroids = scaled_centroids(padded_dim as u32); + let q = make_query(padded_dim, 42); + let (per_vec, interleaved) = make_codes(dim_half, 7); + + let mut lut = vec![0u8; padded_dim * 16]; + let _params = fastscan::build_quantized_lut(&q, ¢roids, &mut lut); + + group.bench_with_input( + BenchmarkId::new("block_scalar", padded_dim), + &padded_dim, + |bench, _| { + let mut results = [0u16; 32]; + bench.iter(|| { + fastscan::fastscan_block_scalar( + black_box(&interleaved), + black_box(&lut), + dim_half, + &mut results, + ); + black_box(results[0]) + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("block_dispatch", padded_dim), + &padded_dim, + |bench, _| { + let scan = fastscan::fastscan_dispatch().scan_block; + let mut results = [0u16; 32]; + bench.iter(|| { + scan( + black_box(&interleaved), + black_box(&lut), + dim_half, + &mut results, + ); + black_box(results[0]) + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("tq_adc_x32", padded_dim), + &padded_dim, + |bench, _| { + bench.iter(|| { + let mut acc = 0.0f32; + for cv in per_vec.iter() { + acc += tq_l2_adc_scaled(black_box(&q), black_box(cv), 1.0, ¢roids); + } + black_box(acc) + }); + }, + ); + } + group.finish(); +} + +fn bench_build_lut(c: &mut Criterion) { + let mut group = c.benchmark_group("fastscan_build_lut"); + + for &padded_dim in PADDED_DIMS { + let centroids = scaled_centroids(padded_dim as u32); + let q = make_query(padded_dim, 42); + let mut lut = vec![0u8; padded_dim * 16]; + + group.bench_with_input( + BenchmarkId::new("build_quantized_lut", padded_dim), + &padded_dim, + |bench, _| { + bench.iter(|| { + let params = fastscan::build_quantized_lut(black_box(&q), ¢roids, &mut lut); + black_box(params.scale) + }); + }, + ); + } + group.finish(); +} + +/// End-to-end mutable-segment scan: FastScan-filtered MVCC path vs the plain +/// per-candidate loop (`brute_force_search_filtered`, same TQ-ADC estimator). +fn bench_mutable_scan(c: &mut Criterion) { + use moon::vector::segment::mutable::MutableSegment; + use moon::vector::turbo_quant::collection::{ + BuildMode, CollectionMetadata, QuantizationConfig, + }; + use moon::vector::types::DistanceMetric; + use std::sync::Arc; + + moon::vector::distance::init(); + + let mut group = c.benchmark_group("mutable_scan_topk"); + group.sample_size(30); + + for &(dim, n) in &[(128usize, 20_000u32), (768, 5_000)] { + let col = Arc::new(CollectionMetadata::with_build_mode( + 1, + dim as u32, + DistanceMetric::L2, + QuantizationConfig::TurboQuant4, + 42, + BuildMode::Light, + )); + let seg = MutableSegment::new(dim as u32, col); + let mut s = 7u32; + let mut v = vec![0.0f32; dim]; + for i in 0..n { + for x in v.iter_mut() { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + *x = (s as f32) / (u32::MAX as f32) * 2.0 - 1.0; + } + seg.append(i as u64, &v, i as u64 + 1); + } + let query = make_query(dim, 42); + let committed = roaring::RoaringTreemap::new(); + + group.bench_with_input( + BenchmarkId::new(format!("fastscan_mvcc_d{dim}"), n), + &n, + |bench, _| { + bench.iter(|| { + let r = seg.brute_force_search_mvcc( + black_box(&query), + None, + 10, + None, + u64::MAX, + 0, + &committed, + 0, + n as usize, + ); + black_box(r.len()) + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new(format!("plain_adc_d{dim}"), n), + &n, + |bench, _| { + bench.iter(|| { + let r = seg.brute_force_search_filtered(black_box(&query), None, 10, None); + black_box(r.len()) + }); + }, + ); + } + group.finish(); +} + +criterion_group!( + benches, + bench_fastscan_block, + bench_build_lut, + bench_mutable_scan +); +criterion_main!(benches); diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index d61626ee..5dafc55d 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -130,6 +130,19 @@ pinned/dedicated cores). | `--segment-cold-min-qps` | `0.1` | Minimum QPS threshold; segments below this are COLD candidates | | `--vec-diskann-beam-width` | `8` | DiskANN beam width for disk-resident vector search | | `--vec-diskann-cache-levels` | `3` | HNSW upper levels cached in memory for DiskANN hybrid search | +| `--vector-ef-runtime` | `0` (auto) | Default `EF_RUNTIME` for indexes created without one (10-4096; `0` = per-query auto heuristic). Per-index `FT.CONFIG SET EF_RUNTIME` overrides at runtime | +| `--vector-rerank-mult` | `4` | Default `RERANK_MULT` for new indexes: exact-rerank depth, top `mult×k` beam candidates re-scored with true f16 distances (1-64). Per-index `FT.CONFIG` overrides | +| `--vector-exact-beam` | off | Default `EXACT_BEAM` for new indexes: HNSW beam navigates exact f16 sidecar distances (recall ≈ graph-limited; QPS cost grows with dimension). Per-index `FT.CONFIG` overrides | + +## Graph + +| Flag | Default | Description | +|------|---------|-------------| +| `--graph-merge-max-segments` | `8` | Immutable CSR segments per graph before autovacuum merges them | +| `--graph-dead-edge-trigger` | `0.20` | Dead-edge fraction that triggers a segment merge early | +| `--graph-timeout-ms` | `30000` | Default Cypher traversal timeout (0 = unlimited; per-query `TIMEOUT` overrides) | +| `--graph-result-cache-entries` | `256` | Cypher result-cache capacity per graph (repeated read-only queries served from cache) | +| `--graph-result-cache-bytes` | `4194304` | Cypher result-cache byte budget per graph (LRU eviction on either limit) | ## Environment Variables diff --git a/docs/guides/tuning.md b/docs/guides/tuning.md index 0329a79d..ee5789fd 100644 --- a/docs/guides/tuning.md +++ b/docs/guides/tuning.md @@ -373,13 +373,42 @@ single-purpose benchmark rigs. ## Vector search (FT.*) -- `EF_RUNTIME` (per index) trades recall for QPS at query time. +- `EF_RUNTIME` (per index) trades recall for QPS at query time — and is now + **runtime-tunable**: `FT.CONFIG SET EF_RUNTIME ` (10–4096, `0` = auto) + applies to the next `FT.SEARCH` immediately, persists across restarts, and needs no + index rebuild. Use it to walk the recall/QPS curve on a live index (e.g. drop ef + during traffic spikes, raise it for offline evaluation). +- **Chasing the last recall points (→ 1.0)?** Two more runtime knobs, both per index, + both persisted, both applied on the next query: + - `FT.CONFIG SET RERANK_MULT ` (1–64, default 4) deepens the exact-rerank + stage: the top `n·k` beam candidates are re-scored with true f16 distances before + truncation. Cheap (`~n·k·dim` f16 decodes per segment) and recovers true neighbors + the quantized ADC ranking dropped just below the default 4·k cut. Try 8–16 first. + - `FT.CONFIG SET EXACT_BEAM ON` goes further: the HNSW beam itself navigates + with exact f16 distances instead of quantized estimates, so recall becomes + graph-limited (~Qdrant parity at equal ef) rather than quantization-limited. QPS + cost grows with dimension (an f16 row is ~4× the bytes of a TQ4 code); benchmark + at your dim before enabling fleet-wide. Segments without an exact-rerank sidecar + (pre-HQ-1 disk reloads) silently keep the quantized beam. + - Escalation order at a fixed recall target: raise `EF_RUNTIME` → raise + `RERANK_MULT` → `EXACT_BEAM ON`. Each step costs more QPS than the one before. + - Fleet-wide starting values: `--vector-ef-runtime`, `--vector-rerank-mult`, and + `--vector-exact-beam` (CLI or `moon.conf`) set the defaults every NEW index is + created with, so recall-sensitive deployments don't have to re-issue FT.CONFIG + per index. Per-index `FT.CONFIG SET` always overrides the server default. - Set `COMPACT_THRESHOLD` at or above your expected dataset size if you want a single final compaction; explicit `FT.COMPACT` on a small mutable segment is a no-op below the threshold. -- Match quantization to dimension: **SQ8** (or full-precision HNSW) for ≤ 384-d - embeddings; **TQ4** shines at 768-d and above. Validate recall with real embeddings, - not random vectors. +- Match quantization to dimension **and metric**: **SQ8** (or full-precision HNSW) for + ≤ 384-d embeddings; **TQ4** shines at 768-d and above and is strongest on the + unit-sphere metrics (COSINE / IP). TQ on raw L2 uses a norm-corrected distance + estimator (an earlier norm-scaled ranking collapsed on unnormalized data, which is + why L2 indexes default to SQ8); SQ8 remains the recommended choice for raw-L2 + workloads. Validate recall with real embeddings, not random vectors. +- **Query cost scales with segment count**: each FT.SEARCH runs the full ef beam on + *every* graph segment on *every* shard (cost ≈ shards × segments × ef). An index + that accumulated 50+ segments during a bulk load answers the same query 4–5× slower + than the same index merged to 1 segment per shard. See the settle recipe below. - Details: [Vector search guide](../vector-search-guide.md). ## Vector bulk load and compaction @@ -418,3 +447,25 @@ knob only in these cases: fires per shard, so bulk loads parallelize across shards automatically. Co-locate related vectors with hash tags only if you also do multi-key KV ops on them; vector search itself scatter-gathers across shards regardless. +- **`--max-unflushed-immutable-segments 0` during million-scale bulk loads.** The + write-stall guard (default 20) counts **total** immutable segments, not just + unflushed ones. A 1M+ load accumulates segments faster than background merges retire + them, so the guard trips permanently and every write returns `MOONERR busy` — + measured 24× ingest slowdown (4,500 → 190 vec/s) on otherwise-idle hardware. Disable + it for the load, restore the default for steady-state serving (it exists to bound + memory under pathological churn). Loaders must still retry on `MOONERR busy` — it is + backpressure, not an error. +- **Settle before latency-sensitive serving.** After a bulk load, merge each shard's + segments to one: `VACUUM VECTOR ` force-merges all immutable segments + (recall-gated GraphUnion). Two caveats: (1) over the wire it acts on the + **connection's local shard only** — with SO_REUSEPORT per-shard listeners, issue it + over ~8× more fresh connections than shards and repeat until `FT.INFO + graph_segments` stops shrinking; (2) merging is currently slow at scale (~1 h for + 1.18M × 200-d on 8 ARM vCPUs, largely single-threaded). The payoff on that dataset: + same index, same ef, **4–5× higher QPS** (e.g. 888 vs 336 qps at recall 0.83). The + index serves correct results throughout — settle when you can, not before you must. +- **Merged indexes need higher ef for the same recall ceiling.** Multi-segment search + unions independent per-segment beams, so an unmerged index over-scans and reaches + higher recall at a given ef (0.9865 vs 0.933 at ef=256 on glove-1.18M). After + settling, raise `FT.CONFIG SET EF_RUNTIME` (e.g. 512) to reclaim the >0.95 + band — still several× faster than the unmerged equivalent. diff --git a/docs/vector-search-guide.md b/docs/vector-search-guide.md index 8f4e10ce..1b4242e7 100644 --- a/docs/vector-search-guide.md +++ b/docs/vector-search-guide.md @@ -46,11 +46,18 @@ FT.CREATE ON HASH PREFIX ... | `DISTANCE_METRIC` | L2 | L2, COSINE, IP | Distance function | | `M` | 16 | 2-64 | HNSW max neighbors per layer. Higher = better recall, more memory | | `EF_CONSTRUCTION` | 200 | 10-4096 | HNSW build effort. Higher = better graph quality, slower compaction | -| `EF_RUNTIME` | auto | 10-4096 | Search beam width. 0/omit = auto: max(k×15, 200). Higher = better recall, lower QPS | +| `EF_RUNTIME` | auto | 10-4096 | Search beam width. 0/omit = auto: max(k×15, 200). Higher = better recall, lower QPS. Tunable at runtime: `FT.CONFIG SET EF_RUNTIME ` (0 = restore auto) — applies to the next search, no rebuild | | `COMPACT_THRESHOLD` | 1000 | 100-100000 | Min vectors before auto-compaction. Higher = fewer larger HNSW graphs | -| `QUANTIZATION` | TQ4 | TQ1-TQ4, SQ8 | Compression level. TQ4 = 4-bit (best compression), SQ8 = 8-bit (higher recall) | +| `QUANTIZATION` | TQ4 (COSINE/IP), SQ8 (L2) | TQ1-TQ4, SQ8 | Compression level. TQ4 = 4-bit (best compression, strongest on unit-sphere metrics; L2 uses a norm-corrected estimator), SQ8 = 8-bit (higher recall, all metrics — the default for L2) | | `BUILD_MODE` | LIGHT | LIGHT, EXACT | HNSW build quality vs resource trade-off (see below) | +Two more per-index recall knobs are runtime-only (set via `FT.CONFIG`, persisted, applied on the next search): + +| FT.CONFIG parameter | Default | Range | Description | +|---------------------|---------|-------|-------------| +| `RERANK_MULT` | 4 | 1-64 | Exact-rerank depth: re-score the top `mult×k` beam candidates with true f16 distances before truncation. Deeper recovers neighbors the quantized ranking dropped, at ~`mult·k·dim` f16 decodes per segment | +| `EXACT_BEAM` | OFF | ON/OFF | Navigate the HNSW beam with exact f16 distances instead of quantized estimates — recall becomes graph-limited (~1.0 at high ef). QPS cost grows with dimension; segments without an exact-rerank sidecar keep the quantized beam | + ### BUILD_MODE: Light vs Exact | Aspect | LIGHT (default) | EXACT | diff --git a/src/command/vector_search/ft_config.rs b/src/command/vector_search/ft_config.rs index 943bbced..a65b5115 100644 --- a/src/command/vector_search/ft_config.rs +++ b/src/command/vector_search/ft_config.rs @@ -165,6 +165,75 @@ fn ft_config_set( } idx.merge_recall_tolerance = parsed; Frame::SimpleString(Bytes::from_static(b"OK")) + } else if param.eq_ignore_ascii_case(b"EF_RUNTIME") { + // Runtime-tunable search beam width, same range as FT.CREATE's + // EF_RUNTIME (10-4096). 0 restores the auto heuristic (max(k*20, 200) + // with dimension boost). Query path reads meta.hnsw_ef_runtime per + // search, so the change applies to the next FT.SEARCH immediately. + let parsed: u32 = match std::str::from_utf8(value) + .ok() + .and_then(|s| s.parse::().ok()) + { + Some(v) if v == 0 || (10..=4096).contains(&v) => v, + Some(_) => { + return Frame::Error(Bytes::from_static( + b"ERR EF_RUNTIME must be 10-4096 (or 0 for auto)", + )); + } + None => { + return Frame::Error(Bytes::from_static(b"ERR invalid EF_RUNTIME value")); + } + }; + idx.meta.hnsw_ef_runtime = parsed; + // Persist so the value survives a server restart. + store.save_index_meta_sidecar(); + Frame::SimpleString(Bytes::from_static(b"OK")) + } else if param.eq_ignore_ascii_case(b"RERANK_MULT") { + // HQ-1 depth: re-score the top `mult · k` beam candidates against the + // f16 exact sidecar. Raising it recovers true neighbors the quantized + // ADC ranking pushed below the 4·k default (the main lever for the + // last few recall points below EXACT_BEAM); cost ~mult·k·dim f16 + // decodes per segment. Applies to the next FT.SEARCH immediately. + let parsed: u32 = match std::str::from_utf8(value) + .ok() + .and_then(|s| s.parse::().ok()) + { + Some(v) if (1..=64).contains(&v) => v, + Some(_) => { + return Frame::Error(Bytes::from_static(b"ERR RERANK_MULT must be 1-64")); + } + None => { + return Frame::Error(Bytes::from_static(b"ERR invalid RERANK_MULT value")); + } + }; + idx.meta.rerank_mult = parsed; + // Persist so the value survives a server restart. + store.save_index_meta_sidecar(); + Frame::SimpleString(Bytes::from_static(b"OK")) + } else if param.eq_ignore_ascii_case(b"EXACT_BEAM") { + // Navigate the HNSW beam with exact f16 sidecar distances instead of + // quantized ADC: recall becomes graph-limited (~1.0 at ef 256) at a + // QPS cost that grows with dimension. Segments without a sidecar + // (pre-HQ-1 reloads) silently keep the ADC beam. + let parsed = if value.eq_ignore_ascii_case(b"ON") + || value == b"1" + || value.eq_ignore_ascii_case(b"TRUE") + { + true + } else if value.eq_ignore_ascii_case(b"OFF") + || value == b"0" + || value.eq_ignore_ascii_case(b"FALSE") + { + false + } else { + return Frame::Error(Bytes::from_static( + b"ERR EXACT_BEAM value must be ON or OFF", + )); + }; + idx.meta.exact_beam = parsed; + // Persist so the value survives a server restart. + store.save_index_meta_sidecar(); + Frame::SimpleString(Bytes::from_static(b"OK")) } else { Frame::Error(Bytes::from_static(b"ERR unknown config parameter")) } @@ -211,6 +280,19 @@ fn ft_config_get( use std::fmt::Write as _; let _ = write!(buf, "{}", idx.merge_recall_tolerance); Frame::BulkString(Bytes::from(buf)) + } else if param.eq_ignore_ascii_case(b"EF_RUNTIME") { + let mut buf = String::with_capacity(8); + use std::fmt::Write as _; + let _ = write!(buf, "{}", idx.meta.hnsw_ef_runtime); + Frame::BulkString(Bytes::from(buf)) + } else if param.eq_ignore_ascii_case(b"RERANK_MULT") { + let mut buf = String::with_capacity(8); + use std::fmt::Write as _; + let _ = write!(buf, "{}", idx.meta.rerank_mult); + Frame::BulkString(Bytes::from(buf)) + } else if param.eq_ignore_ascii_case(b"EXACT_BEAM") { + let val = if idx.meta.exact_beam { "ON" } else { "OFF" }; + Frame::BulkString(Bytes::from(val)) } else { Frame::Error(Bytes::from_static(b"ERR unknown config parameter")) } diff --git a/src/command/vector_search/ft_create.rs b/src/command/vector_search/ft_create.rs index 1282b20c..03d18355 100644 --- a/src/command/vector_search/ft_create.rs +++ b/src/command/vector_search/ft_create.rs @@ -433,6 +433,18 @@ pub fn ft_create( }); } + // Server-wide starting values for the tuning knobs + // (--vector-ef-runtime / --vector-rerank-mult / --vector-exact-beam). + // An explicit EF_RUNTIME in FT.CREATE wins over the server default; + // RERANK_MULT / EXACT_BEAM have no FT.CREATE syntax, so they always + // start at the server default and are tuned per index via FT.CONFIG. + let create_defaults = crate::vector::store::vector_create_defaults(); + let effective_ef_runtime = if first_hnsw_ef_runtime == 0 { + create_defaults.ef_runtime + } else { + first_hnsw_ef_runtime + }; + // Build IndexMeta from the first (default) field for backward compatibility let default_field = &vector_fields[0]; let meta = IndexMeta { @@ -442,7 +454,7 @@ pub fn ft_create( metric: default_field.metric, hnsw_m: first_hnsw_m, hnsw_ef_construction: first_hnsw_ef_construction, - hnsw_ef_runtime: first_hnsw_ef_runtime, + hnsw_ef_runtime: effective_ef_runtime, compact_threshold: first_compact_threshold, source_field: default_field.field_name.clone(), key_prefixes: prefixes.clone(), @@ -453,6 +465,8 @@ pub fn ft_create( merge_mode: first_merge_mode, keep_raw: first_keep_raw, db_index, + rerank_mult: create_defaults.rerank_mult, + exact_beam: create_defaults.exact_beam, }; let index_name_clone = meta.name.clone(); @@ -548,6 +562,7 @@ fn parse_vector_field_params(args: &[Frame], pos: &mut usize) -> Result Result Result 0 { idx.meta.hnsw_ef_runtime as usize } else { @@ -130,6 +134,7 @@ pub(super) fn search_local_raw( dirty_set: &[], dimension: dim as u32, ef_defaulted, + tuning, }; let results = idx.segments.search_mvcc( &query_f32, @@ -156,6 +161,7 @@ pub(super) fn search_local_raw( dirty_set: &[], dimension: dim as u32, ef_defaulted, + tuning, }; let results = fs.segments.search_mvcc( &query_f32, @@ -284,6 +290,10 @@ pub fn search_local_filtered( // AE-1: remember whether ef came from the heuristic (segment-level // adaptive-ef estimates only apply then, never over a user EF_RUNTIME). let ef_defaulted = idx.meta.hnsw_ef_runtime == 0; + let tuning = crate::vector::types::SearchTuning { + rerank_mult: idx.meta.rerank_mult, + exact_beam: idx.meta.exact_beam, + }; let ef_search = if idx.meta.hnsw_ef_runtime > 0 { idx.meta.hnsw_ef_runtime as usize } else { @@ -313,6 +323,7 @@ pub fn search_local_filtered( dirty_set: &[], dimension: dim as u32, ef_defaulted, + tuning, }; let results = idx.segments.search_mvcc( &query_f32, @@ -335,6 +346,7 @@ pub fn search_local_filtered( dirty_set: &[], dimension: dim as u32, ef_defaulted, + tuning, }; let results = fs.segments.search_mvcc( &query_f32, diff --git a/src/command/vector_search/hybrid.rs b/src/command/vector_search/hybrid.rs index 8579b379..66f768d8 100644 --- a/src/command/vector_search/hybrid.rs +++ b/src/command/vector_search/hybrid.rs @@ -593,6 +593,10 @@ pub(super) fn run_dense_knn( dirty_set: &[], dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning { + rerank_mult: idx.meta.rerank_mult, + exact_beam: idx.meta.exact_beam, + }, }; idx.segments .search_mvcc(&query_f32, k, ef_search, &mut idx.scratch, None, &mvcc_ctx) @@ -615,6 +619,10 @@ pub(super) fn run_dense_knn( dirty_set: &[], dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning { + rerank_mult: idx.meta.rerank_mult, + exact_beam: idx.meta.exact_beam, + }, }; fs.segments .search_mvcc(&query_f32, k, ef_search, &mut fs.scratch, None, &mvcc_ctx) diff --git a/src/command/vector_search/recommend.rs b/src/command/vector_search/recommend.rs index 0c9d507c..09848b7b 100644 --- a/src/command/vector_search/recommend.rs +++ b/src/command/vector_search/recommend.rs @@ -348,6 +348,10 @@ pub fn ft_recommend( dirty_set: &[], dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning { + rerank_mult: idx.meta.rerank_mult, + exact_beam: idx.meta.exact_beam, + }, }; // Dispatch to correct field segments diff --git a/src/command/vector_search/tests.rs b/src/command/vector_search/tests.rs index fe1f3cab..acd88eba 100644 --- a/src/command/vector_search/tests.rs +++ b/src/command/vector_search/tests.rs @@ -1403,8 +1403,346 @@ fn test_parse_ft_search_args_without_limit() { assert_eq!(count, usize::MAX); } +#[test] +fn test_ft_create_l2_defaults_to_sq8() { + let _metrics_guard = METRICS_LOCK.read(); + use crate::vector::turbo_quant::collection::QuantizationConfig; + + // L2 + no explicit QUANTIZATION → SQ8 (TQ's norm-scaled ADC collapses on + // unnormalized L2 data; recall < 0.01 measured on gist-960). + let mut store = VectorStore::new(); + ft_create( + &mut store, + &mut crate::text::store::TextStore::new(), + &ft_create_args(), + 0, + ); + #[allow(clippy::unwrap_used)] // index just created + let idx = store.get_index_mut(b"myidx").unwrap(); + assert_eq!(idx.meta.quantization, QuantizationConfig::Sq8); + + // COSINE + no explicit QUANTIZATION → TQ4 default unchanged. + let mut args = ft_create_args(); + args[0] = bulk(b"cosidx"); + let l2_pos = args + .iter() + .position(|f| matches!(f, Frame::BulkString(b) if b.as_ref() == b"L2")) + .expect("L2 arg present"); + args[l2_pos] = bulk(b"COSINE"); + let mut store2 = VectorStore::new(); + ft_create( + &mut store2, + &mut crate::text::store::TextStore::new(), + &args, + 0, + ); + #[allow(clippy::unwrap_used)] // index just created + let idx2 = store2.get_index_mut(b"cosidx").unwrap(); + assert_eq!(idx2.meta.quantization, QuantizationConfig::TurboQuant4); + + // Explicit TQ4 + L2 is honored (warns, not rejected). + let mut args3 = ft_create_args(); + args3[0] = bulk(b"tq4idx"); + let cnt_pos = args3 + .iter() + .position(|f| matches!(f, Frame::BulkString(b) if b.as_ref() == b"6")) + .expect("param count present"); + args3[cnt_pos] = bulk(b"8"); + args3.push(bulk(b"QUANTIZATION")); + args3.push(bulk(b"TQ4")); + let mut store3 = VectorStore::new(); + ft_create( + &mut store3, + &mut crate::text::store::TextStore::new(), + &args3, + 0, + ); + #[allow(clippy::unwrap_used)] // index just created + let idx3 = store3.get_index_mut(b"tq4idx").unwrap(); + assert_eq!(idx3.meta.quantization, QuantizationConfig::TurboQuant4); +} + // -- FT.CONFIG tests -- +#[test] +fn test_ft_config_ef_runtime_set_get() { + let _metrics_guard = METRICS_LOCK.read(); + let mut store = VectorStore::new(); + let args = ft_create_args(); + ft_create( + &mut store, + &mut crate::text::store::TextStore::new(), + &args, + 0, + ); + + // SET a pinned beam width; the query path reads meta.hnsw_ef_runtime live. + let set_args = vec![ + bulk(b"SET"), + bulk(b"myidx"), + bulk(b"EF_RUNTIME"), + bulk(b"64"), + ]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!(matches!(result, Frame::SimpleString(_)), "{result:?}"); + #[allow(clippy::unwrap_used)] // index just created above + let idx = store.get_index_mut(b"myidx").unwrap(); + assert_eq!(idx.meta.hnsw_ef_runtime, 64); + + // GET reflects the new value. + let get_args = vec![bulk(b"GET"), bulk(b"myidx"), bulk(b"EF_RUNTIME")]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &get_args, + 0, + ); + match &result { + Frame::BulkString(b) => assert_eq!(&b[..], b"64"), + other => panic!("expected BulkString 64, got {other:?}"), + } + + // 0 restores the auto heuristic. + let set_args = vec![ + bulk(b"SET"), + bulk(b"myidx"), + bulk(b"EF_RUNTIME"), + bulk(b"0"), + ]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!(matches!(result, Frame::SimpleString(_))); + #[allow(clippy::unwrap_used)] // index exists + let idx = store.get_index_mut(b"myidx").unwrap(); + assert_eq!(idx.meta.hnsw_ef_runtime, 0); + + // Out-of-range and non-numeric values are rejected. + for bad in [&b"5"[..], b"5000", b"abc"] { + let set_args = vec![bulk(b"SET"), bulk(b"myidx"), bulk(b"EF_RUNTIME"), bulk(bad)]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!(matches!(result, Frame::Error(_)), "{bad:?} -> {result:?}"); + } +} + +#[test] +fn test_ft_config_rerank_mult_set_get() { + let _metrics_guard = METRICS_LOCK.read(); + let mut store = VectorStore::new(); + let args = ft_create_args(); + ft_create( + &mut store, + &mut crate::text::store::TextStore::new(), + &args, + 0, + ); + + // Default is 4 (the HQ-1 baseline oversample). + let get_args = vec![bulk(b"GET"), bulk(b"myidx"), bulk(b"RERANK_MULT")]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &get_args, + 0, + ); + match &result { + Frame::BulkString(b) => assert_eq!(&b[..], b"4"), + other => panic!("expected BulkString 4, got {other:?}"), + } + + // SET a deeper oversample; the query path reads meta.rerank_mult live. + let set_args = vec![ + bulk(b"SET"), + bulk(b"myidx"), + bulk(b"RERANK_MULT"), + bulk(b"16"), + ]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!(matches!(result, Frame::SimpleString(_)), "{result:?}"); + #[allow(clippy::unwrap_used)] // index just created above + let idx = store.get_index_mut(b"myidx").unwrap(); + assert_eq!(idx.meta.rerank_mult, 16); + + let get_args = vec![bulk(b"GET"), bulk(b"myidx"), bulk(b"RERANK_MULT")]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &get_args, + 0, + ); + match &result { + Frame::BulkString(b) => assert_eq!(&b[..], b"16"), + other => panic!("expected BulkString 16, got {other:?}"), + } + + // Out-of-range and non-numeric values are rejected. + for bad in [&b"0"[..], b"65", b"abc"] { + let set_args = vec![ + bulk(b"SET"), + bulk(b"myidx"), + bulk(b"RERANK_MULT"), + bulk(bad), + ]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!(matches!(result, Frame::Error(_)), "{bad:?} -> {result:?}"); + } +} + +#[test] +fn test_ft_config_tuning_knobs_are_db_scoped() { + // WS5a interplay: an index created in db 1 is invisible to FT.CONFIG + // from any other db — RERANK_MULT / EXACT_BEAM can only be set (and + // read) through the owning db. + let _metrics_guard = METRICS_LOCK.read(); + let mut store = VectorStore::new(); + let args = ft_create_args(); + ft_create( + &mut store, + &mut crate::text::store::TextStore::new(), + &args, + 1, + ); + + // SET from db 0 (wrong db) must not resolve the index. + let set_args = vec![ + bulk(b"SET"), + bulk(b"myidx"), + bulk(b"RERANK_MULT"), + bulk(b"32"), + ]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!(matches!(result, Frame::Error(_)), "{result:?}"); + + // SET from db 1 (owning db) succeeds; knob lands on the db-1 index. + for (param, val) in [(&b"RERANK_MULT"[..], &b"32"[..]), (b"EXACT_BEAM", b"ON")] { + let set_args = vec![bulk(b"SET"), bulk(b"myidx"), bulk(param), bulk(val)]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 1, + ); + assert!(matches!(result, Frame::SimpleString(_)), "{result:?}"); + } + #[allow(clippy::unwrap_used)] // created above + let idx_db1 = store.get_index_mut_for_db(b"myidx", 1).unwrap(); + assert_eq!(idx_db1.meta.rerank_mult, 32); + assert!(idx_db1.meta.exact_beam); + + // GET is db-scoped too: wrong db errors, owning db reads the knob. + let get_args = vec![bulk(b"GET"), bulk(b"myidx"), bulk(b"EXACT_BEAM")]; + let r0 = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &get_args, + 0, + ); + assert!(matches!(r0, Frame::Error(_)), "{r0:?}"); + let r1 = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &get_args, + 1, + ); + match &r1 { + Frame::BulkString(b) => assert_eq!(&b[..], b"ON"), + other => panic!("expected BulkString ON, got {other:?}"), + } +} + +#[test] +fn test_ft_config_exact_beam_on_off() { + let _metrics_guard = METRICS_LOCK.read(); + let mut store = VectorStore::new(); + let args = ft_create_args(); + ft_create( + &mut store, + &mut crate::text::store::TextStore::new(), + &args, + 0, + ); + + // Default is OFF. + let get_args = vec![bulk(b"GET"), bulk(b"myidx"), bulk(b"EXACT_BEAM")]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &get_args, + 0, + ); + match &result { + Frame::BulkString(b) => assert_eq!(&b[..], b"OFF"), + other => panic!("expected BulkString OFF, got {other:?}"), + } + + for (val, expect) in [ + (&b"ON"[..], true), + (b"OFF", false), + (b"1", true), + (b"0", false), + (b"true", true), + (b"false", false), + ] { + let set_args = vec![bulk(b"SET"), bulk(b"myidx"), bulk(b"EXACT_BEAM"), bulk(val)]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!( + matches!(result, Frame::SimpleString(_)), + "{val:?} -> {result:?}" + ); + #[allow(clippy::unwrap_used)] // index exists + let idx = store.get_index_mut(b"myidx").unwrap(); + assert_eq!(idx.meta.exact_beam, expect, "value {val:?}"); + } + + // Invalid values are rejected. + let set_args = vec![ + bulk(b"SET"), + bulk(b"myidx"), + bulk(b"EXACT_BEAM"), + bulk(b"maybe"), + ]; + let result = ft_config( + &mut store, + &mut crate::text::store::TextStore::new(), + &set_args, + 0, + ); + assert!(matches!(result, Frame::Error(_)), "{result:?}"); +} + #[test] fn test_ft_config_autocompact_on_off() { let _metrics_guard = METRICS_LOCK.read(); diff --git a/src/config.rs b/src/config.rs index 613bc3ac..da2083e7 100644 --- a/src/config.rs +++ b/src/config.rs @@ -684,6 +684,55 @@ pub struct ServerConfig { #[arg(long = "graph-timeout-ms", default_value_t = 30_000)] pub graph_timeout_ms: u64, + /// Cypher result-cache capacity: maximum cached query results per graph. + /// + /// The cache serves repeated read-only Cypher queries without + /// re-executing them (doorkeeper-gated admission; invalidated on any + /// write to the graph). Raising it helps dashboards that cycle through + /// many distinct read queries; each graph gets its own cache. + /// + /// Default: 256 entries. + #[arg(long = "graph-result-cache-entries", default_value_t = 256)] + pub graph_result_cache_entries: usize, + + /// Cypher result-cache memory ceiling in bytes (per graph). + /// + /// Entries are evicted LRU when either this byte budget or + /// `--graph-result-cache-entries` is exceeded. + /// + /// Default: 4194304 (4 MiB). + #[arg(long = "graph-result-cache-bytes", default_value_t = 4_194_304)] + pub graph_result_cache_bytes: usize, + + // ── Vector search-tuning defaults (FT.CREATE / FT.CONFIG initial values) ── + /// Default HNSW search beam width (EF_RUNTIME) for indexes created + /// without an explicit `EF_RUNTIME` in FT.CREATE. `0` keeps the per-query + /// auto heuristic (max(k*20, 200) with dimension boost). Range: 10-4096. + /// + /// Per-index `FT.CONFIG SET EF_RUNTIME` still overrides this at + /// runtime; the flag only sets the starting value for NEW indexes. + #[arg(long = "vector-ef-runtime", default_value_t = 0)] + pub vector_ef_runtime: u32, + + /// Default exact-rerank depth multiplier (RERANK_MULT) for new vector + /// indexes: the top `mult × k` beam candidates are re-scored with true + /// f16 sidecar distances before top-k truncation. Range: 1-64. + /// + /// Per-index `FT.CONFIG SET RERANK_MULT` still overrides this at + /// runtime; the flag only sets the starting value for NEW indexes. + #[arg(long = "vector-rerank-mult", default_value_t = 4)] + pub vector_rerank_mult: u32, + + /// Default EXACT_BEAM state for new vector indexes: when set, the HNSW + /// beam navigates with exact f16 sidecar distances instead of quantized + /// ADC estimates — recall becomes graph-limited at a QPS cost that grows + /// with dimension. + /// + /// Per-index `FT.CONFIG SET EXACT_BEAM OFF` still overrides this + /// at runtime; the flag only sets the starting value for NEW indexes. + #[arg(long = "vector-exact-beam", default_value_t = false)] + pub vector_exact_beam: bool, + // ── MA4: Weighted compaction scheduling ─────────────────────────────── /// Minimum seconds before a stale entity is forced to be scheduled by the /// autovacuum daemon regardless of its compaction weight (anti-starvation cap). @@ -836,6 +885,38 @@ impl ServerConfig { /// `Frame`/`anyhow::Error` — kept plain so both `main.rs`'s early-boot /// path and unit tests can format/assert on it directly) when the /// configured count exceeds `MAX_DATABASES`. + /// Validate the vector/graph tuning-default flags (startup error, not a + /// silent clamp — a typo in a fleet-wide default must be loud). Mirrors + /// the per-index FT.CONFIG ranges exactly so a value accepted here is + /// accepted there and vice versa. + pub fn validate_tuning_defaults(&self) -> Result<(), String> { + if self.vector_ef_runtime != 0 && !(10..=4096).contains(&self.vector_ef_runtime) { + return Err(format!( + "--vector-ef-runtime {} out of range (10-4096, or 0 for the auto heuristic)", + self.vector_ef_runtime + )); + } + if !(1..=64).contains(&self.vector_rerank_mult) { + return Err(format!( + "--vector-rerank-mult {} out of range (1-64)", + self.vector_rerank_mult + )); + } + if self.graph_result_cache_entries == 0 { + return Err( + "--graph-result-cache-entries must be >= 1 (the cache cannot be sized to zero entries; it is admission-gated, not disableable by size)" + .to_string(), + ); + } + if self.graph_result_cache_bytes < 4096 { + return Err(format!( + "--graph-result-cache-bytes {} too small (minimum 4096)", + self.graph_result_cache_bytes + )); + } + Ok(()) + } + pub fn validate_databases_bound(&self) -> Result<(), String> { if self.databases > MAX_DATABASES { return Err(format!( @@ -1610,6 +1691,57 @@ mod tests { assert_eq!(config.shards, 0); } + // ── Vector/graph tuning-default flags ───────────────────────────────── + + #[test] + fn test_tuning_default_flags_parse_and_default() { + let config = ServerConfig::parse_from::<[&str; 0], &str>([]); + assert_eq!(config.vector_ef_runtime, 0); + assert_eq!(config.vector_rerank_mult, 4); + assert!(!config.vector_exact_beam); + assert_eq!(config.graph_result_cache_entries, 256); + assert_eq!(config.graph_result_cache_bytes, 4_194_304); + assert!(config.validate_tuning_defaults().is_ok()); + + let config = ServerConfig::parse_from([ + "moon", + "--vector-ef-runtime", + "256", + "--vector-rerank-mult", + "16", + "--vector-exact-beam", + "--graph-result-cache-entries", + "1024", + "--graph-result-cache-bytes", + "8388608", + ]); + assert_eq!(config.vector_ef_runtime, 256); + assert_eq!(config.vector_rerank_mult, 16); + assert!(config.vector_exact_beam); + assert_eq!(config.graph_result_cache_entries, 1024); + assert_eq!(config.graph_result_cache_bytes, 8_388_608); + assert!(config.validate_tuning_defaults().is_ok()); + } + + #[test] + fn test_tuning_default_flags_out_of_range_rejected() { + // Same ranges as FT.CONFIG: ef 10-4096 (or 0), mult 1-64. + for args in [ + vec!["moon", "--vector-ef-runtime", "5"], + vec!["moon", "--vector-ef-runtime", "5000"], + vec!["moon", "--vector-rerank-mult", "0"], + vec!["moon", "--vector-rerank-mult", "65"], + vec!["moon", "--graph-result-cache-entries", "0"], + vec!["moon", "--graph-result-cache-bytes", "100"], + ] { + let config = ServerConfig::parse_from(args.clone()); + assert!( + config.validate_tuning_defaults().is_err(), + "expected startup rejection for {args:?}" + ); + } + } + // ── WS5a round 2 (adversarial review finding 2): --databases bound ──── #[test] diff --git a/src/config/conf_file.rs b/src/config/conf_file.rs index c3765c39..dabd0007 100644 --- a/src/config/conf_file.rs +++ b/src/config/conf_file.rs @@ -103,6 +103,7 @@ pub const BOOL_FLAGS: &[&str] = &[ "unsafe-multishard-aof", "experimental-per-shard-rewrite", "cluster-enabled", + "vector-exact-beam", ]; /// Parse a conf file and return a list of synthesised argv tokens. @@ -456,6 +457,23 @@ mod tests { assert_eq!(tokens, vec!["--cluster-enabled"]); } + #[test] + fn vector_exact_beam_is_a_bool_flag() { + // moon.conf: `vector-exact-beam yes` must synthesise the bare flag. + let tokens = parse_conf_contents( + "vector-exact-beam yes +", + ) + .unwrap(); + assert_eq!(tokens, vec!["--vector-exact-beam"]); + let tokens = parse_conf_contents( + "vector-exact-beam no +", + ) + .unwrap(); + assert!(tokens.is_empty()); + } + #[test] fn multiple_bool_flags() { let contents = "cluster-enabled yes\nconsole-auth-required no\nunsafe-multishard-aof yes\n"; diff --git a/src/graph/cypher/result_cache.rs b/src/graph/cypher/result_cache.rs index e4c161d2..184734d3 100644 --- a/src/graph/cypher/result_cache.rs +++ b/src/graph/cypher/result_cache.rs @@ -298,6 +298,26 @@ pub const DEFAULT_MAX_ENTRIES: usize = 256; /// `DEFAULT_MAX_ENTRIES` entries at both protocol versions with headroom. pub const DEFAULT_MAX_BYTES: usize = 4 * 1024 * 1024; +/// Operator-configured cache limits (`--graph-result-cache-entries` / +/// `--graph-result-cache-bytes`), installed once at startup. `GraphStore` +/// construction reads these; the compiled-in defaults apply when the server +/// never installed any (unit tests, library embedders). +static CONFIGURED_LIMITS: std::sync::OnceLock<(usize, usize)> = std::sync::OnceLock::new(); + +/// Install the server-wide result-cache limits. First write wins; later +/// calls are no-ops (safe for embedded/test servers booting in-process). +pub fn set_configured_limits(max_entries: usize, max_bytes: usize) { + let _ = CONFIGURED_LIMITS.set((max_entries, max_bytes)); +} + +/// `(max_entries, max_bytes)` for new per-graph caches. +pub fn configured_limits() -> (usize, usize) { + CONFIGURED_LIMITS + .get() + .copied() + .unwrap_or((DEFAULT_MAX_ENTRIES, DEFAULT_MAX_BYTES)) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/graph/store.rs b/src/graph/store.rs index a00c492f..8cc51524 100644 --- a/src/graph/store.rs +++ b/src/graph/store.rs @@ -381,12 +381,13 @@ impl GraphStore { )), key_to_node: HashMap::new(), write_gen: 0, - result_cache: parking_lot::Mutex::new( - crate::graph::cypher::result_cache::ResultCache::new( - crate::graph::cypher::result_cache::DEFAULT_MAX_ENTRIES, - crate::graph::cypher::result_cache::DEFAULT_MAX_BYTES, - ), - ), + result_cache: parking_lot::Mutex::new({ + // Operator-configured limits (--graph-result-cache-*), + // falling back to the compiled-in defaults. + let (max_entries, max_bytes) = + crate::graph::cypher::result_cache::configured_limits(); + crate::graph::cypher::result_cache::ResultCache::new(max_entries, max_bytes) + }), }, ); // Bump version AFTER successful graph creation (monotonicity-on-success diff --git a/src/main.rs b/src/main.rs index 29435903..2019c7f6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -224,6 +224,23 @@ fn main() -> anyhow::Result<()> { return Err(anyhow::anyhow!(msg)); } + // Vector/graph tuning defaults: validate the ranges loudly, then install + // the process-wide starting values (read by FT.CREATE and GraphStore + // construction; per-index FT.CONFIG still overrides at runtime). + if let Err(msg) = config.validate_tuning_defaults() { + return Err(anyhow::anyhow!(msg)); + } + moon::vector::store::set_vector_create_defaults(moon::vector::store::VectorCreateDefaults { + ef_runtime: config.vector_ef_runtime, + rerank_mult: config.vector_rerank_mult, + exact_beam: config.vector_exact_beam, + }); + #[cfg(feature = "graph")] + moon::graph::cypher::result_cache::set_configured_limits( + config.graph_result_cache_entries, + config.graph_result_cache_bytes, + ); + // Non-jemalloc builds: warn if operator explicitly set --memory-arenas-cap #[cfg(not(feature = "jemalloc"))] if config.memory_arenas_cap != 8 { diff --git a/src/server/conn/handler_monoio/ft.rs b/src/server/conn/handler_monoio/ft.rs index 8f25757d..eda64704 100644 --- a/src/server/conn/handler_monoio/ft.rs +++ b/src/server/conn/handler_monoio/ft.rs @@ -325,15 +325,40 @@ pub(super) async fn try_handle_ft_command( return true; } if cmd.eq_ignore_ascii_case(b"FT.CONFIG") { + // SET must broadcast: each shard holds its own index partition, and a + // local-only write would leave 7/8 shards on the old value (the tokio + // sharded handler already broadcasts FT.CONFIG via its catch-all). + // GET stays local — post-broadcast, every shard agrees. The + // connection's SELECTed db rides the broadcast (WS5a db scoping). let db_index = conn.selected_db as u8; - let response = crate::shard::slice::with_shard(|s| { - crate::command::vector_search::ft_config( - &mut s.vector_store, - &mut s.text_store, - cmd_args, + let is_set = cmd_args + .first() + .and_then(|f| match f { + crate::protocol::Frame::BulkString(b) => Some(b.as_ref()), + _ => None, + }) + .is_some_and(|s| s.eq_ignore_ascii_case(b"SET")); + let response = if is_set { + crate::shard::coordinator::broadcast_vector_command( + std::sync::Arc::new(frame.clone()), + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, db_index, ) - }); + .await + } else { + crate::shard::slice::with_shard(|s| { + crate::command::vector_search::ft_config( + &mut s.vector_store, + &mut s.text_store, + cmd_args, + db_index, + ) + }) + }; responses.push(response); return true; } diff --git a/src/server/embedded.rs b/src/server/embedded.rs index 8ba543ea..c33b9395 100644 --- a/src/server/embedded.rs +++ b/src/server/embedded.rs @@ -92,6 +92,23 @@ pub async fn run_embedded( // cwd data) before validating it — matches the binary entry in main.rs. config.resolve_dir(); + // Install the vector/graph tuning defaults exactly like the binary entry + // (first write wins — a host process that already set them keeps its + // values). Validation errors are startup errors here too. + if let Err(msg) = config.validate_tuning_defaults() { + anyhow::bail!(msg); + } + crate::vector::store::set_vector_create_defaults(crate::vector::store::VectorCreateDefaults { + ef_runtime: config.vector_ef_runtime, + rerank_mult: config.vector_rerank_mult, + exact_beam: config.vector_exact_beam, + }); + #[cfg(feature = "graph")] + crate::graph::cypher::result_cache::set_configured_limits( + config.graph_result_cache_entries, + config.graph_result_cache_bytes, + ); + // Validate / create persistence directory up front. std::fs::create_dir_all(&config.dir).with_context(|| { format!( diff --git a/src/vector/distance/fastscan.rs b/src/vector/distance/fastscan.rs index 23a5b50e..1128ac71 100644 --- a/src/vector/distance/fastscan.rs +++ b/src/vector/distance/fastscan.rs @@ -1,8 +1,25 @@ -//! VPSHUFB FastScan distance kernel for IVF posting list scanning. +//! FastScan distance kernels: register-resident 4-bit LUT accumulation. //! -//! Computes approximate distances for 32 vectors simultaneously using -//! precomputed u8 LUT lookups. The AVX2 path uses VPSHUFB (_mm256_shuffle_epi8) -//! for 32 parallel table lookups per instruction. +//! Implements the FAISS FastScan technique (André et al., VLDB'15): TQ4 codes +//! are stored in interleaved 32-vector blocks and approximate distances for +//! all 32 vectors are accumulated simultaneously via in-register SIMD table +//! lookups — VPSHUFB (`_mm256_shuffle_epi8`) on x86_64, TBL (`vqtbl1q_u8`) on +//! aarch64. Each coordinate's 16-entry u8 LUT fits exactly one SIMD register, +//! so the inner loop never touches RAM for table lookups. +//! +//! Distance semantics (all kernels are bit-identical): +//! - per sub-dim, the two nibble distances are widened to u16 and summed +//! exactly (max 255 + 255 = 510, no wrap-around), +//! - the running per-vector accumulator is u16 with *saturating* addition, +//! so pathological high-dim/high-scale combinations clamp at 65535 instead +//! of wrapping (a clamped distance ranks last, which is the correct failure +//! mode for a far-away candidate). +//! +//! LUT construction ([`build_quantized_lut`]) follows the FAISS scheme: a +//! per-coordinate bias (the minimum distance for that coordinate) is +//! subtracted before quantization and a single global scale maps the residual +//! range to u8, chosen so that the worst-case accumulated sum also fits u16. +//! The true float distance is reconstructed as `acc / scale + bias`. //! //! The scalar fallback produces identical results on all architectures. @@ -23,7 +40,8 @@ static FASTSCAN_DISPATCH: OnceLock = OnceLock::new(); /// Initialize the FastScan dispatch table. /// -/// Selects AVX2 kernel on x86_64 when available, scalar otherwise. +/// Kernel selection: NEON TBL on aarch64 (baseline feature, no runtime +/// detection needed), AVX2 VPSHUFB on x86_64 when available, scalar otherwise. /// Safe to call multiple times (OnceLock guarantees single init). pub fn init_fastscan() { FASTSCAN_DISPATCH.get_or_init(|| { @@ -39,7 +57,15 @@ pub fn init_fastscan() { } } - // Scalar fallback for all platforms. + #[cfg(target_arch = "aarch64")] + { + return FastScanDispatch { + scan_block: fastscan_block_neon, + }; + } + + // Scalar fallback for all other platforms. + #[allow(unreachable_code)] FastScanDispatch { scan_block: fastscan_block_scalar, } @@ -72,7 +98,7 @@ pub fn fastscan_dispatch() -> &'static FastScanDispatch { /// represents a pair of coordinates. /// `results`: Output accumulated u16 distances for 32 vectors (caller-provided). /// -/// No allocations. +/// Accumulation is saturating (see module docs). No allocations. pub fn fastscan_block_scalar(codes: &[u8], lut: &[u8], dim_half: usize, results: &mut [u16; 32]) { // Zero-initialize results. *results = [0u16; 32]; @@ -89,7 +115,7 @@ pub fn fastscan_block_scalar(codes: &[u8], lut: &[u8], dim_half: usize, results: let dist_lo = lut[lut_lo_base + lo_idx] as u16; let dist_hi = lut[lut_hi_base + hi_idx] as u16; - results[v] += dist_lo + dist_hi; + results[v] = results[v].saturating_add(dist_lo + dist_hi); } } } @@ -101,7 +127,9 @@ pub fn fastscan_block_scalar(codes: &[u8], lut: &[u8], dim_half: usize, results: /// 1. Load 32 nibble-packed bytes -> split lo/hi nibbles /// 2. Broadcast 16-byte LUT to both lanes of __m256i /// 3. VPSHUFB: 32 parallel lookups for even and odd coordinates -/// 4. Accumulate into u16 accumulators (zero-extend u8 -> u16 to avoid overflow) +/// 4. Widen each lookup result to u16 BEFORE adding (u8 + u8 can exceed 255), +/// then accumulate with saturating u16 adds — bit-identical to the scalar +/// kernel for the full u8 LUT range. /// /// # Safety /// Caller must verify AVX2 is available via `is_x86_feature_detected!("avx2")`. @@ -120,7 +148,8 @@ pub unsafe fn fastscan_block_avx2( let lo_mask = _mm256_set1_epi8(0x0F); let zero = _mm256_setzero_si256(); - // Two u16 accumulators: acc_lo holds vectors 0..15, acc_hi holds vectors 16..31. + // Two u16 accumulators covering 32 vectors (lane-interleaved layout, + // untangled at store time). let mut acc_lo = _mm256_setzero_si256(); // 16 x u16 let mut acc_hi = _mm256_setzero_si256(); // 16 x u16 @@ -150,60 +179,201 @@ pub unsafe fn fastscan_block_avx2( let dist_lo = _mm256_shuffle_epi8(lut_lo_vec, lo_nibbles); let dist_hi = _mm256_shuffle_epi8(lut_hi_vec, hi_nibbles); - // Add lo + hi distances (u8 + u8, still fits u8 for individual coord pair). - // Then widen to u16 and accumulate. - let dist_sum = _mm256_add_epi8(dist_lo, dist_hi); + // Widen each u8 lookup to u16 first, then add exactly (max 510), + // then accumulate with saturating u16 adds. Adding as u8 would wrap + // for LUT pairs summing past 255. + let lo_w0 = _mm256_unpacklo_epi8(dist_lo, zero); + let lo_w1 = _mm256_unpackhi_epi8(dist_lo, zero); + let hi_w0 = _mm256_unpacklo_epi8(dist_hi, zero); + let hi_w1 = _mm256_unpackhi_epi8(dist_hi, zero); - // Zero-extend lower 16 bytes to u16 and accumulate. - let lo_16 = _mm256_unpacklo_epi8(dist_sum, zero); - let hi_16 = _mm256_unpackhi_epi8(dist_sum, zero); + let sum_w0 = _mm256_add_epi16(lo_w0, hi_w0); + let sum_w1 = _mm256_add_epi16(lo_w1, hi_w1); - acc_lo = _mm256_add_epi16(acc_lo, lo_16); - acc_hi = _mm256_add_epi16(acc_hi, hi_16); + acc_lo = _mm256_adds_epu16(acc_lo, sum_w0); + acc_hi = _mm256_adds_epu16(acc_hi, sum_w1); } // Store accumulators to results. - // unpacklo/unpackhi interleaves within 128-bit lanes, so the layout is: - // acc_lo: [v0,v1,v2,v3,v4,v5,v6,v7 | v16,v17,v18,v19,v20,v21,v22,v23] (u16) - // acc_hi: [v8,v9,v10,v11,v12,v13,v14,v15 | v24,v25,v26,v27,v28,v29,v30,v31] (u16) - // We need to extract and interleave properly. - // - // Actually, _mm256_unpacklo_epi8 interleaves bytes from the lower half of each - // 128-bit lane. For 32 input bytes [b0..b31], after unpacklo with zero: - // result = [b0,0,b1,0,...,b7,0 | b16,0,b17,0,...,b23,0] - // And unpackhi: - // result = [b8,0,b9,0,...,b15,0 | b24,0,b25,0,...,b31,0] - // - // So we store and rearrange. + // _mm256_unpacklo_epi8 interleaves bytes within each 128-bit lane, so: + // acc_lo: [v0..v7 | v16..v23] (u16 lanes) + // acc_hi: [v8..v15 | v24..v31] (u16 lanes) let mut tmp_lo = [0u16; 16]; let mut tmp_hi = [0u16; 16]; _mm256_storeu_si256(tmp_lo.as_mut_ptr() as *mut __m256i, acc_lo); _mm256_storeu_si256(tmp_hi.as_mut_ptr() as *mut __m256i, acc_hi); - // Rearrange from lane-interleaved to linear order. - // acc_lo lane 0 (indices 0..7): vectors 0,1,2,3,4,5,6,7 - // acc_lo lane 1 (indices 8..15): vectors 16,17,18,19,20,21,22,23 - // acc_hi lane 0 (indices 0..7): vectors 8,9,10,11,12,13,14,15 - // acc_hi lane 1 (indices 8..15): vectors 24,25,26,27,28,29,30,31 results[0..8].copy_from_slice(&tmp_lo[0..8]); results[8..16].copy_from_slice(&tmp_hi[0..8]); results[16..24].copy_from_slice(&tmp_lo[8..16]); results[24..32].copy_from_slice(&tmp_hi[8..16]); } +/// NEON TBL FastScan: compute distances for 32 vectors in one interleaved block. +/// +/// Each coordinate's 16-entry u8 LUT fits exactly one 128-bit `q` register; +/// `vqtbl1q_u8` performs 16 parallel lookups per instruction. A 32-vector +/// block is processed as two 16-vector halves. Widening uses `vaddl_u8` / +/// `vaddl_high_u8` (exact u8+u8 -> u16), accumulation uses `vqaddq_u16` +/// (saturating) — bit-identical to the scalar kernel. +/// +/// NEON is a baseline feature on aarch64 — no runtime detection required +/// (same convention as `src/vector/distance/neon.rs`). +#[cfg(target_arch = "aarch64")] +pub fn fastscan_block_neon(codes: &[u8], lut: &[u8], dim_half: usize, results: &mut [u16; 32]) { + use std::arch::aarch64::*; + + // Runtime-enforced (release builds included): the unsafe block below does + // raw-pointer loads and must never trust caller discipline alone. Cost is + // two predicted branches per 32-vector block — noise vs the scan itself. + assert!(codes.len() >= dim_half * BLOCK_SIZE); + assert!(lut.len() >= dim_half * 32); + + // SAFETY: NEON is mandatory baseline on aarch64 (target_arch cfg above). + // Per iteration we read 16 bytes at codes[d*32 + half*16] and at + // lut[(2d)*16] / lut[(2d+1)*16], all bounded by the two asserts above. + unsafe { + let nibble_mask = vdupq_n_u8(0x0F); + // Four u16x8 accumulators in linear vector order: v0..7, v8..15, v16..23, v24..31. + let mut acc0 = vdupq_n_u16(0); + let mut acc1 = vdupq_n_u16(0); + let mut acc2 = vdupq_n_u16(0); + let mut acc3 = vdupq_n_u16(0); + + for d in 0..dim_half { + let code_ptr = codes.as_ptr().add(d * BLOCK_SIZE); + let lut_lo = vld1q_u8(lut.as_ptr().add((2 * d) * 16)); + let lut_hi = vld1q_u8(lut.as_ptr().add((2 * d + 1) * 16)); + + // First half: vectors 0..15. + let packed_a = vld1q_u8(code_ptr); + let lo_a = vandq_u8(packed_a, nibble_mask); + let hi_a = vshrq_n_u8(packed_a, 4); + let dist_lo_a = vqtbl1q_u8(lut_lo, lo_a); + let dist_hi_a = vqtbl1q_u8(lut_hi, hi_a); + acc0 = vqaddq_u16( + acc0, + vaddl_u8(vget_low_u8(dist_lo_a), vget_low_u8(dist_hi_a)), + ); + acc1 = vqaddq_u16(acc1, vaddl_high_u8(dist_lo_a, dist_hi_a)); + + // Second half: vectors 16..31. + let packed_b = vld1q_u8(code_ptr.add(16)); + let lo_b = vandq_u8(packed_b, nibble_mask); + let hi_b = vshrq_n_u8(packed_b, 4); + let dist_lo_b = vqtbl1q_u8(lut_lo, lo_b); + let dist_hi_b = vqtbl1q_u8(lut_hi, hi_b); + acc2 = vqaddq_u16( + acc2, + vaddl_u8(vget_low_u8(dist_lo_b), vget_low_u8(dist_hi_b)), + ); + acc3 = vqaddq_u16(acc3, vaddl_high_u8(dist_lo_b, dist_hi_b)); + } + + vst1q_u16(results.as_mut_ptr(), acc0); + vst1q_u16(results.as_mut_ptr().add(8), acc1); + vst1q_u16(results.as_mut_ptr().add(16), acc2); + vst1q_u16(results.as_mut_ptr().add(24), acc3); + } +} + +/// Reconstruction parameters returned by [`build_quantized_lut`]. +/// +/// The true float ADC distance is approximately `acc as f32 / scale + bias` +/// where `acc` is the u16 accumulator produced by a FastScan block kernel. +#[derive(Debug, Clone, Copy)] +pub struct LutParams { + /// Global scale factor applied to (dist - coord_min) before u8 quantization. + pub scale: f32, + /// Sum of per-coordinate minimum distances (subtracted during quantization). + pub bias: f32, +} + +/// Build a quantized u8 distance LUT from a rotated query (FAISS scheme). +/// +/// For each coordinate `c` in `0..padded_dim` the 16 candidate distances +/// `(q_rotated[c] - centroids[k])^2` are computed; the per-coordinate minimum +/// is subtracted (its sum across coordinates becomes `bias`) and the residual +/// is scaled by a single global factor into u8. The scale is the largest +/// value that (a) keeps every entry <= 255 and (b) keeps the worst-case +/// accumulated sum <= u16::MAX, so kernel saturation cannot trigger for +/// in-range data. +/// +/// `centroids` must be the same dimension-scaled codebook used to encode the +/// stored TQ codes (`scaled_centroids(padded_dim)`), NOT the legacy +/// 1/sqrt(768) table — asymmetric codebooks silently destroy recall. +/// +/// `lut_out` must have length >= `q_rotated.len() * 16`. No allocations. +pub fn build_quantized_lut( + q_rotated: &[f32], + centroids: &[f32; 16], + lut_out: &mut [u8], +) -> LutParams { + let padded_dim = q_rotated.len(); + debug_assert!(lut_out.len() >= padded_dim * 16); + + // Pass 1: per-coordinate min/max to derive bias and the global scale. + let mut bias = 0.0f32; + let mut sum_range = 0.0f32; + let mut max_range = 0.0f32; + for &q_val in q_rotated.iter() { + let mut coord_min = f32::MAX; + let mut coord_max = 0.0f32; + for &c in centroids.iter() { + let diff = q_val - c; + let d = diff * diff; + coord_min = coord_min.min(d); + coord_max = coord_max.max(d); + } + bias += coord_min; + let range = coord_max - coord_min; + sum_range += range; + max_range = max_range.max(range); + } + + let scale = if max_range > 0.0 { + (255.0 / max_range).min(65535.0 / sum_range.max(f32::MIN_POSITIVE)) + } else { + 1.0 + }; + + // Pass 2: quantize. Recomputing the 16 distances per coordinate is cheaper + // than allocating a padded_dim*16 f32 scratch on the query path. + for (coord, &q_val) in q_rotated.iter().enumerate() { + let mut coord_min = f32::MAX; + let mut dists = [0.0f32; 16]; + for (k, &c) in centroids.iter().enumerate() { + let diff = q_val - c; + let d = diff * diff; + dists[k] = d; + coord_min = coord_min.min(d); + } + let base = coord * 16; + for (k, &d) in dists.iter().enumerate() { + let scaled = (d - coord_min) * scale + 0.5; + lut_out[base + k] = if scaled >= 255.0 { 255 } else { scaled as u8 }; + } + } + + LutParams { scale, bias } +} + /// Scan all blocks in a posting list and collect top-k results. /// /// `codes`: Full interleaved code buffer from PostingList. -/// `lut`: Precomputed u8 distance LUT (padded_dim * 16 entries). +/// `lut`: Quantized u8 distance LUT (padded_dim * 16 entries). +/// `params`: Scale/bias from [`build_quantized_lut`] for f32 reconstruction. /// `dim_half`: padded_dim / 2. /// `ids`: Vector IDs from PostingList. /// `norms`: Precomputed norms from PostingList. /// `count`: Number of vectors in the posting list. /// `k`: Number of results to keep. /// `results`: Output buffer for SearchResults (caller-provided SmallVec). +#[allow(clippy::too_many_arguments)] pub fn scan_posting_list( codes: &[u8], lut: &[u8], + params: LutParams, dim_half: usize, ids: &[u32], norms: &[f32], @@ -213,8 +383,9 @@ pub fn scan_posting_list( ) { let dispatch = fastscan_dispatch(); let n = count as usize; - let n_blocks = (n + BLOCK_SIZE - 1) / BLOCK_SIZE; + let n_blocks = n.div_ceil(BLOCK_SIZE); let block_bytes = dim_half * BLOCK_SIZE; + let inv_scale = 1.0 / params.scale; let mut block_dists = [0u16; 32]; @@ -230,14 +401,12 @@ pub fn scan_posting_list( &mut block_dists, ); - // Convert u16 quantized distances to f32 and push results. + // Reconstruct approximate f32 ADC distances and push results. for v in 0..vecs_in_block { let global_idx = vec_start + v; let norm = norms[global_idx]; - // Scale back: u16 distance is sum of quantized per-coord distances. - // The actual L2 distance is approximately: norm^2 * (raw_dist / LUT_SCALE_TOTAL) - // For ranking purposes, raw u16 distance * norm^2 preserves ordering. - let dist_f32 = block_dists[v] as f32 * norm * norm; + let adc = block_dists[v] as f32 * inv_scale + params.bias; + let dist_f32 = adc * norm * norm; results.push(SearchResult::new(dist_f32, VectorId(ids[global_idx]))); } } @@ -279,6 +448,24 @@ mod tests { (codes, lut, dim_half) } + /// Deterministic full-range pseudo-random codes + LUT (no 0x7F masking — + /// exercises the widen-before-add and saturating-accumulate paths). + fn make_random_block(dim_half: usize, seed: u32) -> (Vec, Vec) { + let padded_dim = dim_half * 2; + let mut codes = vec![0u8; dim_half * BLOCK_SIZE]; + let mut lut = vec![0u8; padded_dim * 16]; + let mut s = seed; + for b in codes.iter_mut() { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + *b = (s >> 24) as u8; + } + for b in lut.iter_mut() { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + *b = (s >> 24) as u8; // FULL 0..=255 range + } + (codes, lut) + } + #[test] fn test_fastscan_block_scalar_known_distances() { let dim_half = 2; @@ -290,12 +477,9 @@ mod tests { // For vector v: each sub-dim contributes lut[lo_idx] + lut[hi_idx]. // lo_idx = hi_idx = v % 16. lut[coord * 16 + k] = k. - // So per sub-dim: v + v = 2*v. Over dim_half=2 sub-dims: 2 * 2*v = 4*v. - // Wait: we have 2 coordinates per sub-dim (even + odd). // dist_lo = lut[(2*d) * 16 + lo_idx] = lo_idx = v // dist_hi = lut[(2*d+1) * 16 + hi_idx] = hi_idx = v - // Per sub-dim: v + v = 2*v. - // Over dim_half=2: 2 * 2*v = 4*v. + // Per sub-dim: v + v = 2*v. Over dim_half=2: 2 * 2*v = 4*v. for v in 0..n_vectors { assert_eq!( results[v], @@ -360,6 +544,22 @@ mod tests { } } + #[test] + fn test_fastscan_scalar_saturates_instead_of_wrapping() { + // All-255 LUT with enough sub-dims to exceed u16::MAX: 200 sub-dims + // * 510 per sub-dim = 102000 > 65535. Accumulator must clamp, not wrap. + let dim_half = 200; + let padded_dim = dim_half * 2; + let codes = vec![0xFFu8; dim_half * BLOCK_SIZE]; // all nibbles = 15 + let lut = vec![255u8; padded_dim * 16]; + + let mut results = [0u16; 32]; + fastscan_block_scalar(&codes, &lut, dim_half, &mut results); + for v in 0..BLOCK_SIZE { + assert_eq!(results[v], u16::MAX, "v={v} must saturate at u16::MAX"); + } + } + #[test] fn test_scan_posting_list_scalar_topk() { init_fastscan(); @@ -394,6 +594,10 @@ mod tests { scan_posting_list( &codes, &lut, + LutParams { + scale: 1.0, + bias: 0.0, + }, dim_half, &ids, &norms, @@ -414,47 +618,172 @@ mod tests { #[cfg(target_arch = "x86_64")] #[test] - fn test_fastscan_block_avx2_matches_scalar() { + fn test_fastscan_block_avx2_matches_scalar_full_range() { if !is_x86_feature_detected!("avx2") { return; } - // Test with random-ish data. - let dim_half = 64; // 128 coordinates - let padded_dim = dim_half * 2; - let mut codes = vec![0u8; dim_half * BLOCK_SIZE]; - let mut lut = vec![0u8; padded_dim * 16]; + // Full-range LUT (0..=255): locks in the widen-before-add fix — the + // old kernel wrapped u8 pairs summing past 255. + for (dim_half, seed) in [(1usize, 7u32), (3, 11), (64, 42), (200, 99)] { + let (codes, lut) = make_random_block(dim_half, seed); - // Fill with deterministic pseudo-random data. - let mut s = 42u32; - for b in codes.iter_mut() { - s = s.wrapping_mul(1664525).wrapping_add(1013904223); - *b = (s >> 24) as u8; + let mut scalar_results = [0u16; 32]; + fastscan_block_scalar(&codes, &lut, dim_half, &mut scalar_results); + + let mut avx2_results = [0u16; 32]; + // SAFETY: AVX2 checked above. + unsafe { + fastscan_block_avx2(&codes, &lut, dim_half, &mut avx2_results); + } + + assert_eq!( + avx2_results, scalar_results, + "AVX2 vs scalar mismatch (dim_half={dim_half})" + ); } - for b in lut.iter_mut() { + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn test_fastscan_block_neon_matches_scalar_full_range() { + // Full-range LUT (0..=255) across several dims, including a + // saturating case (dim_half=200 with high LUT values). + for (dim_half, seed) in [(1usize, 7u32), (3, 11), (64, 42), (200, 99)] { + let (codes, lut) = make_random_block(dim_half, seed); + + let mut scalar_results = [0u16; 32]; + fastscan_block_scalar(&codes, &lut, dim_half, &mut scalar_results); + + let mut neon_results = [0u16; 32]; + fastscan_block_neon(&codes, &lut, dim_half, &mut neon_results); + + assert_eq!( + neon_results, scalar_results, + "NEON vs scalar mismatch (dim_half={dim_half})" + ); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn test_fastscan_block_neon_saturates() { + let dim_half = 200; + let codes = vec![0xFFu8; dim_half * BLOCK_SIZE]; + let lut = vec![255u8; dim_half * 2 * 16]; + let mut results = [0u16; 32]; + fastscan_block_neon(&codes, &lut, dim_half, &mut results); + for v in 0..BLOCK_SIZE { + assert_eq!(results[v], u16::MAX, "v={v} must saturate at u16::MAX"); + } + } + + #[test] + fn test_build_quantized_lut_reconstruction_accuracy() { + use crate::vector::turbo_quant::codebook::scaled_centroids; + + // Reconstructed FastScan distance must track the exact f32 ADC + // distance closely for a realistic rotated-unit-vector magnitude. + let padded_dim = 128usize; + let dim_half = padded_dim / 2; + let centroids = scaled_centroids(padded_dim as u32); + let sigma = 1.0 / (padded_dim as f32).sqrt(); + + // Deterministic query: coordinates in +/- 2*sigma (typical for FWHT output). + let mut s = 12345u32; + let mut q = vec![0.0f32; padded_dim]; + for v in q.iter_mut() { s = s.wrapping_mul(1664525).wrapping_add(1013904223); - // LUT values must be in [0, 127] to avoid overflow when adding lo+hi as u8. - *b = ((s >> 24) as u8) & 0x7F; + *v = ((s as f32) / (u32::MAX as f32) * 4.0 - 2.0) * sigma; } - let mut scalar_results = [0u16; 32]; - fastscan_block_scalar(&codes, &lut, dim_half, &mut scalar_results); + let mut lut = vec![0u8; padded_dim * 16]; + let params = build_quantized_lut(&q, ¢roids, &mut lut); + assert!(params.scale > 0.0, "scale must be positive"); + + // Random codes for 32 vectors, exact f32 ADC vs reconstructed. + let mut codes_per_vec = vec![vec![0u8; dim_half]; 32]; + for cv in codes_per_vec.iter_mut() { + for b in cv.iter_mut() { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + *b = (s >> 24) as u8; + } + } - let mut avx2_results = [0u16; 32]; - // SAFETY: AVX2 checked above. - unsafe { - fastscan_block_avx2(&codes, &lut, dim_half, &mut avx2_results); + // Interleave manually: out[d*32 + v] = codes_per_vec[v][d]. + let mut interleaved = vec![0u8; dim_half * BLOCK_SIZE]; + for (v, cv) in codes_per_vec.iter().enumerate() { + for (d, &byte) in cv.iter().enumerate() { + interleaved[d * BLOCK_SIZE + v] = byte; + } } - for v in 0..BLOCK_SIZE { - assert_eq!( - avx2_results[v], scalar_results[v], - "AVX2 vs scalar mismatch at v={v}: avx2={}, scalar={}", - avx2_results[v], scalar_results[v] + let mut acc = [0u16; 32]; + fastscan_block_scalar(&interleaved, &lut, dim_half, &mut acc); + + for (v, cv) in codes_per_vec.iter().enumerate() { + // Exact f32 ADC for this code. + let mut exact = 0.0f32; + for (d, &byte) in cv.iter().enumerate() { + let lo = (byte & 0x0F) as usize; + let hi = (byte >> 4) as usize; + let d_lo = q[2 * d] - centroids[lo]; + let d_hi = q[2 * d + 1] - centroids[hi]; + exact += d_lo * d_lo + d_hi * d_hi; + } + + let reconstructed = acc[v] as f32 / params.scale + params.bias; + let err = (reconstructed - exact).abs(); + // Quantization error bound: each coordinate contributes at most + // ~0.5/scale rounding error. + let bound = padded_dim as f32 / params.scale; + assert!( + err <= bound, + "v={v}: reconstructed {reconstructed} vs exact {exact} (err {err} > bound {bound})" ); } } + #[test] + fn test_build_quantized_lut_preserves_ranking() { + use crate::vector::turbo_quant::codebook::scaled_centroids; + + // A code identical to the query's own quantization must rank strictly + // better than a code with every nibble maximally wrong. + let padded_dim = 64usize; + let dim_half = padded_dim / 2; + let centroids = scaled_centroids(padded_dim as u32); + + // Query sitting exactly on centroid 4 in every coordinate. + let q = vec![centroids[4]; padded_dim]; + let mut lut = vec![0u8; padded_dim * 16]; + let params = build_quantized_lut(&q, ¢roids, &mut lut); + + let good_code = vec![0x44u8; dim_half]; // both nibbles = 4 + let bad_code = vec![0xFFu8; dim_half]; // both nibbles = 15 (far centroid) + + let mut interleaved = vec![0u8; dim_half * BLOCK_SIZE]; + for (d, &b) in good_code.iter().enumerate() { + interleaved[d * BLOCK_SIZE] = b; + } + for (d, &b) in bad_code.iter().enumerate() { + interleaved[d * BLOCK_SIZE + 1] = b; + } + + let mut acc = [0u16; 32]; + fastscan_block_scalar(&interleaved, &lut, dim_half, &mut acc); + + let good = acc[0] as f32 / params.scale + params.bias; + let bad = acc[1] as f32 / params.scale + params.bias; + assert!( + good < bad, + "exact-match code ({good}) must rank before far code ({bad})" + ); + // The exact-match code's residual is 0 in every coordinate, so its + // accumulator must be exactly 0 (all its LUT entries are per-coord minima). + assert_eq!(acc[0], 0, "exact-match accumulator must be zero"); + } + #[test] fn test_fastscan_dispatch_init() { init_fastscan(); @@ -477,4 +806,27 @@ mod tests { // Vector 0: dist = lut[0*16+1] + lut[1*16+1] = 1 + 1 = 2 assert_eq!(results[0], 2); } + + #[test] + fn test_fastscan_dispatch_matches_scalar_full_range() { + // Whatever kernel the dispatch selected on this machine must be + // bit-identical to the scalar reference across the full LUT range. + init_fastscan(); + let d = fastscan_dispatch(); + + for (dim_half, seed) in [(3usize, 5u32), (64, 21), (200, 77)] { + let (codes, lut) = make_random_block(dim_half, seed); + + let mut scalar_results = [0u16; 32]; + fastscan_block_scalar(&codes, &lut, dim_half, &mut scalar_results); + + let mut dispatch_results = [0u16; 32]; + (d.scan_block)(&codes, &lut, dim_half, &mut dispatch_results); + + assert_eq!( + dispatch_results, scalar_results, + "dispatch vs scalar mismatch (dim_half={dim_half})" + ); + } + } } diff --git a/src/vector/hnsw/search.rs b/src/vector/hnsw/search.rs index 2a5cbe5b..5faafa22 100644 --- a/src/vector/hnsw/search.rs +++ b/src/vector/hnsw/search.rs @@ -224,6 +224,7 @@ pub fn hnsw_search( None, &[], 0, + None, ) } @@ -253,6 +254,7 @@ pub fn hnsw_search_subcent( None, sub_centroid_signs, sub_sign_bytes_per_vec, + None, ) } @@ -274,6 +276,7 @@ pub fn hnsw_search_filtered( allow_bitmap: Option<&RoaringBitmap>, sub_centroid_signs: &[u8], sub_sign_bpv: usize, + exact_f16: Option<&[u16]>, ) -> SmallVec<[SearchResult; 32]> { let num_nodes = graph.num_nodes(); if num_nodes == 0 { @@ -365,6 +368,64 @@ pub fn hnsw_search_filtered( // Capture immutable slice of rotated query (after mutation phase is done) let q_rotated: &[f32] = scratch.query_rotated.as_slice(); + // TQ's LUT sum is a UNIT-SPHERE distance d̂² between the normalized query + // and the decoded unit direction; scaling by ‖a‖² ranks correctly only for + // unit-sphere metrics (COSINE/IP, where encode normalizes too). For raw L2 + // reconstruct the true metric from the stored document norm and the query + // norm: ‖a−q‖² = (‖a‖−‖q‖)² + ‖a‖·‖q‖·d̂². Without this, small-norm + // vectors rank near regardless of direction (gist-960 recall 0.002). + let l2_adjust = !is_sq8 && collection.metric == DistanceMetric::L2; + let finish_tq = move |sphere_sum: f32, norm: f32| -> f32 { + if l2_adjust { + let diff = norm - q_norm; + diff * diff + norm * q_norm * sphere_sum + } else { + sphere_sum * (norm * norm) + } + }; + // EXACT_BEAM: when `exact_f16` is Some, both distance closures below score + // candidates with TRUE f16 distances from the exact-rerank sidecar instead + // of quantized ADC — beam candidate SELECTION becomes exact, so recall is + // graph-limited rather than quantization-limited. Same distance + // conventions as `ImmutableSegment::rerank_exact` (true squared L2 for + // L2; `2 − 2·⟨q̂,x̂⟩` for the unit-sphere metrics), so cross-segment merge + // ordering stays consistent. Out-of-range sidecar reads (never expected — + // the sidecar is BFS-complete) fall through to the quantized path. + let exact_kernels = distance::table(); + let mut exact_q_unit: SmallVec<[f32; 512]> = SmallVec::new(); + if exact_f16.is_some() && collection.metric != DistanceMetric::L2 { + exact_q_unit = SmallVec::from_slice(query); + let n: f32 = exact_q_unit.iter().map(|x| x * x).sum::().sqrt(); + if n > 0.0 { + let inv = 1.0 / n; + for v in exact_q_unit.iter_mut() { + *v *= inv; + } + } + } + let exact_metric_is_l2 = collection.metric == DistanceMetric::L2; + let exact_q: &[f32] = if exact_metric_is_l2 { + query + } else { + &exact_q_unit + }; + let exact_dist = move |vec_f16: &[u16]| -> f32 { + if exact_metric_is_l2 { + (exact_kernels.f16_l2)(exact_q, vec_f16) + } else { + let (dot, xsq) = (exact_kernels.f16_dot_normsq)(exact_q, vec_f16); + if xsq > 0.0 { + // f16 rounding can push cos slightly outside [-1, 1]; clamp + // so distances stay in the metric's [0, 4] range. + let cos = (dot / xsq.sqrt()).clamp(-1.0, 1.0); + 2.0 - 2.0 * cos + } else { + // Zero vector: normalized form undefined — neutral distance. + 2.0 + } + } + }; + // SQ8 has no shared codebook; use a zero placeholder so codebook_16() never // logs a spurious "empty codebook" error (the SQ8 closures ignore the LUT). let zero_codebook = [0.0f32; 16]; @@ -452,6 +513,13 @@ pub fn hnsw_search_filtered( // Hot path: processes `code_len` bytes (nibble-packed TQ codes) with LUT lookups. // For 384d: code_len ≈ 192, 384 nibble lookups per candidate, called ~500 times per query. let dist_bfs = |bfs_pos: u32| -> f32 { + // EXACT_BEAM override: true f16 distance from the sidecar. + if let Some(raw) = exact_f16 { + let start = bfs_pos as usize * dim; + if let Some(vec_f16) = raw.get(start..start + dim) { + return exact_dist(vec_f16); + } + } let offset = bfs_pos as usize * bytes_per_code; if is_sq8 { // SQ8: decode (min, scale) from the slot trailer, SIMD-dispatched @@ -484,7 +552,6 @@ pub fn hnsw_search_filtered( let code_only = &vectors_tq[offset..offset + code_len]; let norm_bytes = &vectors_tq[offset + code_len..offset + bytes_per_code]; let norm = f32::from_le_bytes([norm_bytes[0], norm_bytes[1], norm_bytes[2], norm_bytes[3]]); - let norm_sq = norm * norm; if use_subcent { // Hot path: 90%+ of search time. Optimization strategy: @@ -556,7 +623,7 @@ pub fn hnsw_search_filtered( s0 += adc_lut[qi * 32 + (byte & 0x0F) as usize * 2 + s_lo]; s1 += adc_lut[(qi + 1) * 32 + (byte >> 4) as usize * 2 + s_hi]; } - ((s0 + s1) + (s2 + s3) + (s4 + s5) + (s6 + s7)) * norm_sq + finish_tq((s0 + s1) + (s2 + s3) + (s4 + s5) + (s6 + s7), norm) } else { // 4-way unrolled with independent accumulators for ILP. // Uses unsafe get_unchecked to eliminate bounds checks in the hot loop. @@ -606,12 +673,23 @@ pub fn hnsw_search_filtered( s0 += adc_lut[qi * 16 + (byte & 0x0F)]; s1 += adc_lut[(qi + 1) * 16 + (byte >> 4)]; } - ((s0 + s1) + (s2 + s3) + (s4 + s5) + (s6 + s7)) * norm_sq + finish_tq((s0 + s1) + (s2 + s3) + (s4 + s5) + (s6 + s7), norm) } }; // LUT-based budgeted distance with early termination. let dist_bfs_budgeted = |bfs_pos: u32, budget: f32| -> f32 { + // EXACT_BEAM override: the exact distance is cheap enough relative to + // its accuracy win that the ADC early-exit budget is not needed — + // compute the true distance and let the caller's domination check + // (`d >= worst_dist`) discard it. + if let Some(raw) = exact_f16 { + let start = bfs_pos as usize * dim; + if let Some(vec_f16) = raw.get(start..start + dim) { + let _ = budget; + return exact_dist(vec_f16); + } + } let offset = bfs_pos as usize * bytes_per_code; if is_sq8 { // SQ8 ADC is cheap and exact; ignore the budget early-exit. @@ -646,9 +724,21 @@ pub fn hnsw_search_filtered( let norm = f32::from_le_bytes([norm_bytes[0], norm_bytes[1], norm_bytes[2], norm_bytes[3]]); let norm_sq = norm * norm; if norm_sq <= 0.0 { - return 0.0; + // Zero-norm vector: sphere distance is undefined. For L2 the true + // distance is exactly ||q||^2; unit-sphere metrics keep legacy 0.0. + return if l2_adjust { q_norm * q_norm } else { 0.0 }; } - let scaled_budget = budget / norm_sq; + // Invert the metric transform so the early-exit check on the raw + // sphere sum stays valid: candidate qualifies iff + // finish_tq(sum, norm) <= budget <=> sum <= scaled_budget. + let scaled_budget = if l2_adjust { + let diff = norm - q_norm; + // May go negative (candidate cannot qualify) -- the first budget + // check then returns MAX, which is correct. + (budget - diff * diff) / (norm * q_norm).max(f32::MIN_POSITIVE) + } else { + budget / norm_sq + }; let mut sum = 0.0f32; let check_interval = 16; let chunks = code_only.len() / check_interval; @@ -747,7 +837,7 @@ pub fn hnsw_search_filtered( sum += adc_lut[(qi + 1) * 16 + (byte >> 4) as usize]; } } - sum * norm_sq + finish_tq(sum, norm) }; // Step 2: Upper layer greedy descent (original node ID space) @@ -1296,6 +1386,114 @@ mod tests { // ── hnsw_search tests ───────────────────────────────────────────── + /// L2-on-TQ metric correction: with UNNORMALIZED vectors (norms 0.5..17), + /// beam search must rank by true L2, not sphere_dist*norm^2 (the gist-960 + /// recall-collapse bug). Ground truth is exact f32 L2. + #[test] + fn test_search_l2_unnormalized_recall() { + distance::init(); + let n = 400usize; + let dim = 32usize; + let collection = CollectionMetadata::new( + 1, + dim as u32, + DistanceMetric::L2, + QuantizationConfig::TurboQuant4, + 42, + ); + let padded = collection.padded_dimension as usize; + let signs_vec: Vec = collection.fwht_sign_flips.as_slice().to_vec(); + let signs = signs_vec.as_slice(); + + // Unit directions scaled to varying norms. + let mut vectors = Vec::with_capacity(n); + let mut codes = Vec::with_capacity(n); + let mut work = vec![0.0f32; padded]; + for i in 0..n { + let mut v = lcg_f32(dim, (i * 7 + 13) as u32); + normalize(&mut v); + let scale = 0.5 + (i % 17) as f32; + v.iter_mut().for_each(|x| *x *= scale); + let boundaries = collection.codebook_boundaries_15(); + let code = encode_tq_mse_scaled(&v, signs, boundaries, &mut work); + vectors.push(v); + codes.push(code); + } + + let bytes_per_code = padded / 2 + 4; + let mut tq_buffer_orig: Vec = Vec::with_capacity(n * bytes_per_code); + for code in &codes { + tq_buffer_orig.extend_from_slice(&code.codes); + tq_buffer_orig.extend_from_slice(&code.norm.to_le_bytes()); + } + + // Build graph with a TRUE-L2 pairwise oracle (exact f32) so the graph + // topology itself is metric-faithful; this test isolates SEARCH-side + // scoring. + let mut builder = HnswBuilder::new(16, 200, 12345); + for i in 0..n { + let a = i; // builder assigns ids in insertion order + let _ = a; + builder.insert(|a: u32, b: u32| { + let va = &vectors[a as usize]; + let vb = &vectors[b as usize]; + va.iter() + .zip(vb.iter()) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + }); + } + let graph = builder.build(bytes_per_code as u32); + + let mut tq_buffer_bfs = vec![0u8; n * bytes_per_code]; + for bfs_pos in 0..n { + let orig_id = graph.to_original(bfs_pos as u32) as usize; + let src = orig_id * bytes_per_code; + let dst = bfs_pos * bytes_per_code; + tq_buffer_bfs[dst..dst + bytes_per_code] + .copy_from_slice(&tq_buffer_orig[src..src + bytes_per_code]); + } + + // Unnormalized query. + let mut query = lcg_f32(dim, 991); + normalize(&mut query); + query.iter_mut().for_each(|x| *x *= 6.5); + + // Exact f32 L2 ground truth. + let k = 10; + let mut exact: Vec<(f32, u32)> = vectors + .iter() + .enumerate() + .map(|(i, v)| { + let d = v + .iter() + .zip(query.iter()) + .map(|(x, y)| (x - y) * (x - y)) + .sum::(); + (d, i as u32) + }) + .collect(); + exact.sort_by(|a, b| a.0.total_cmp(&b.0)); + let want: std::collections::HashSet = exact[..k].iter().map(|(_, i)| *i).collect(); + + let mut scratch = SearchScratch::new(n as u32, padded as u32); + let results = hnsw_search( + &graph, + &tq_buffer_bfs, + &query, + &collection, + k, + 128, + &mut scratch, + ); + let got: std::collections::HashSet = results.iter().map(|r| r.id.0).collect(); + let overlap = want.intersection(&got).count(); + assert!( + overlap >= 6, + "L2 beam recall on unnormalized data: {overlap}/{k}" + ); + } + #[test] fn test_search_empty_graph() { distance::init(); @@ -1470,6 +1668,7 @@ mod tests { None, &[], 0, + None, ); assert_eq!(unfiltered.len(), filtered.len()); @@ -1508,6 +1707,7 @@ mod tests { Some(&bitmap), &[], 0, + None, ); for r in &results { assert!( diff --git a/src/vector/index_persist.rs b/src/vector/index_persist.rs index 2fc6310f..5339b360 100644 --- a/src/vector/index_persist.rs +++ b/src/vector/index_persist.rs @@ -39,7 +39,7 @@ //! [compaction_weight: f32 LE] ← NEW in v3 //! ``` //! -//! ## Format v4 (current — WS5a db-scoped indexes) +//! ## Format v4 (WS5a db-scoped indexes) //! //! Extends v3 with a single `db_index: u8` byte appended after //! `compaction_weight`, tagging which logical db (`SELECT 0..databases-1`) @@ -54,6 +54,17 @@ //! ... (same as v3 fields) ... //! [db_index: u8] ← NEW in v4 //! ``` +//! +//! ## Format v5 (current — search-tuning knobs) +//! +//! Extends v4 with the FT.CONFIG search-tuning knobs appended after +//! `db_index`. v1-v4 sidecars are read with the defaults (mult 4, beam off). +//! +//! ```text +//! Per index: +//! ... (same as v4 fields) ... +//! [rerank_mult: u32 LE] [exact_beam: u8] [reserved: 3B] ← NEW in v5 +//! ``` use std::io::{self, Read, Write}; use std::path::Path; @@ -69,6 +80,9 @@ const VERSION_V1: u8 = 1; const VERSION_V2: u8 = 2; const VERSION_V3: u8 = 3; const VERSION_V4: u8 = 4; +/// v5 appends per-index search-tuning knobs (rerank_mult u32 LE, +/// exact_beam u8, 3 reserved bytes) after the v4 db_index byte. +const VERSION_V5: u8 = 5; /// Default compaction weight used when reading v1/v2 sidecars without a stored weight. const DEFAULT_WEIGHT_ON_LOAD: f32 = 1.0; @@ -100,10 +114,10 @@ fn serialize_index_metas_v1(metas: &[&IndexMeta]) -> Vec { /// from `vector_fields[0]` for backward compatibility), then appends the full /// `vector_fields` array. pub fn serialize_index_metas(metas: &[&IndexMeta]) -> Vec { - // Wrap with default weight=1.0 and delegate to the current (v4) serializer. + // Wrap with default weight=1.0 and delegate to the current (v5) serializer. let pairs: Vec<(&IndexMeta, f32)> = metas.iter().map(|&m| (m, DEFAULT_WEIGHT_ON_LOAD)).collect(); - serialize_index_metas_v4(&pairs) + serialize_index_metas_v5(&pairs) } /// Serialize `(IndexMeta, compaction_weight)` pairs to bytes using v3 format (W3-deep). @@ -116,14 +130,23 @@ fn serialize_index_metas_v3(pairs: &[(&IndexMeta, f32)]) -> Vec { serialize_index_metas_versioned(pairs, VERSION_V3) } -/// Serialize `(IndexMeta, compaction_weight)` pairs to bytes using the -/// current v4 format (WS5a): v3 plus a trailing `db_index: u8` per index. +/// Serialize `(IndexMeta, compaction_weight)` pairs to bytes using the v4 +/// format (WS5a): v3 plus a trailing `db_index: u8` per index. Kept for +/// v4-migration test coverage; production writers use +/// [`serialize_index_metas_v5`] (via [`serialize_index_metas`]). pub fn serialize_index_metas_v4(pairs: &[(&IndexMeta, f32)]) -> Vec { serialize_index_metas_versioned(pairs, VERSION_V4) } -/// Shared v3/v4 serializer — `version` selects whether the trailing -/// `db_index` byte (v4) is written. +/// Serialize `(IndexMeta, compaction_weight)` pairs to bytes using the +/// current v5 format: v4 plus the per-index search-tuning knobs +/// (`rerank_mult` u32 LE, `exact_beam` u8, 3 reserved bytes). +pub fn serialize_index_metas_v5(pairs: &[(&IndexMeta, f32)]) -> Vec { + serialize_index_metas_versioned(pairs, VERSION_V5) +} + +/// Shared v3/v4/v5 serializer — `version` selects which trailing extensions +/// (`db_index` byte at v4, search-tuning knobs at v5) are written. fn serialize_index_metas_versioned(pairs: &[(&IndexMeta, f32)], version: u8) -> Vec { let mut buf = Vec::with_capacity(256); @@ -155,6 +178,13 @@ fn serialize_index_metas_versioned(pairs: &[(&IndexMeta, f32)], version: u8) -> if version >= VERSION_V4 { buf.push(m.db_index); } + + // v5 extension: search-tuning knobs + if version >= VERSION_V5 { + buf.extend_from_slice(&m.rerank_mult.to_le_bytes()); + buf.push(m.exact_beam as u8); + buf.extend_from_slice(&[0u8; 3]); // reserved + } } buf @@ -189,12 +219,14 @@ fn write_v1_per_index(buf: &mut Vec, m: &IndexMeta) { } } -/// Deserialize IndexMeta list from bytes. Handles v1, v2, and v3 formats. +/// Deserialize IndexMeta list from bytes. Handles v1 through v5 formats. /// -/// v1/v2 data is auto-migrated: +/// Older data is auto-migrated: /// - v1: single source_field wrapped into 1-element `vector_fields`. /// - v2: full field array; `compaction_weight` defaults to 1.0. /// - v3: full field array + explicit `compaction_weight` per index. +/// - v4: v3 + `db_index` per index (pre-v4 defaults to 0). +/// - v5: v4 + `rerank_mult`/`exact_beam` (pre-v5 defaults to 4 / OFF). /// /// Returns `(IndexMeta, compaction_weight)` pairs. pub fn deserialize_index_metas_with_weights(data: &[u8]) -> io::Result> { @@ -205,11 +237,7 @@ pub fn deserialize_index_metas_with_weights(data: &[u8]) -> io::Result io::Result= VERSION_V5 { + let mult = read_u32(data, &mut cursor)?; + let beam = read_u8(data, &mut cursor)? != 0; + cursor += 3; // reserved + let mult = if (1..=64).contains(&mult) { mult } else { 4 }; + (mult, beam) + } else { + (4, false) + }; + let meta = IndexMeta { name: meta_base.0, dimension, @@ -304,6 +346,8 @@ pub fn deserialize_index_metas_with_weights(data: &[u8]) -> io::Result>, + /// Per-index recall/QPS knobs (FT.CONFIG RERANK_MULT / EXACT_BEAM) — + /// threaded so pooled == serial identity holds. Immutable segments only; + /// warm segments carry no f16 sidecar and ignore it. + pub tuning: crate::vector::types::SearchTuning, pub reply: flume::Sender>, } @@ -163,10 +167,17 @@ fn run_job(job: &SegmentSearchJob, scratch: &mut SearchScratch) -> SmallVec<[Sea let filter = job.filter.as_deref(); match &job.segment { GraphSegmentRef::Immutable(seg) => match filter { - Some(bm) => { - seg.search_filtered(&job.query, job.fetch_k, job.ef_search, scratch, Some(bm)) + Some(bm) => seg.search_filtered_with_tuning( + &job.query, + job.fetch_k, + job.ef_search, + scratch, + Some(bm), + job.tuning, + ), + None => { + seg.search_with_tuning(&job.query, job.fetch_k, job.ef_search, scratch, job.tuning) } - None => seg.search(&job.query, job.fetch_k, job.ef_search, scratch), }, GraphSegmentRef::Warm(seg) => match filter { Some(bm) => { @@ -290,6 +301,7 @@ mod tests { scratch: SearchScratch::new(0, padded_dimension(dim)), key_hash_to_key: idx.key_hash_to_key.clone(), ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), } } @@ -374,6 +386,7 @@ mod tests { scratch: SearchScratch::new(0, padded_dimension(dim)), key_hash_to_key: key_map.clone(), ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), }; let results = futures::executor::block_on( SegmentHolder::search_mvcc_yielding_with_pool( diff --git a/src/vector/segment/holder.rs b/src/vector/segment/holder.rs index 74ee48a7..95e8b348 100644 --- a/src/vector/segment/holder.rs +++ b/src/vector/segment/holder.rs @@ -46,6 +46,9 @@ pub struct MvccContext<'a> { /// user `EF_RUNTIME`) — saturation-certified segments may then run at /// their min-ef estimate. Always false when the user pinned ef. pub ef_defaulted: bool, + /// Per-index recall/QPS knobs (FT.CONFIG RERANK_MULT / EXACT_BEAM), + /// copied from `IndexMeta` at search entry. + pub tuning: crate::vector::types::SearchTuning, } /// Snapshot of all segments at a point in time. @@ -211,6 +214,9 @@ pub struct SearchSnapshot { /// user `EF_RUNTIME`) — saturation-certified segments may then run at /// their min-ef estimate. pub ef_defaulted: bool, + /// Per-index recall/QPS knobs (FT.CONFIG RERANK_MULT / EXACT_BEAM), + /// copied from `IndexMeta` at capture. + pub tuning: crate::vector::types::SearchTuning, } /// Lock-free segment holder. Searches load() once at query start and hold @@ -591,9 +597,16 @@ impl SegmentHolder { for imm in &snapshot.immutable { let ef_i = seg_ef(imm.suggested_ef()); if filter_bitmap.is_some() { - all.extend(imm.search_filtered(query_f32, k, ef_i, _scratch, filter_bitmap)); + all.extend(imm.search_filtered_with_tuning( + query_f32, + k, + ef_i, + _scratch, + filter_bitmap, + mvcc.tuning, + )); } else { - all.extend(imm.search(query_f32, k, ef_i, _scratch)); + all.extend(imm.search_with_tuning(query_f32, k, ef_i, _scratch, mvcc.tuning)); } } @@ -750,6 +763,7 @@ impl SegmentHolder { let snapshot_lsn = snap.snapshot_lsn; let my_txn_id = snap.my_txn_id; let mutable_len = snap.mutable_len; + let tuning = snap.tuning; // Prepare TurboQuant_prod query state for mutable search (same as sync). let collection = segments.mutable.collection(); @@ -796,20 +810,28 @@ impl SegmentHolder { fetch_k, ef_search: ef_seg, filter: filter_arc.clone(), + tuning, reply: tx.clone(), }; if pool.submit(job) { pending_replies += 1; } else if graph_filter.is_some() { - all.extend(seg.search_filtered( + all.extend(seg.search_filtered_with_tuning( query_f32, fetch_k, ef_seg, &mut snap.scratch, graph_filter, + tuning, )); } else { - let results = seg.search(query_f32, fetch_k, ef_seg, &mut snap.scratch); + let results = seg.search_with_tuning( + query_f32, + fetch_k, + ef_seg, + &mut snap.scratch, + tuning, + ); if post_filter { if let Some(bm) = filter_ref { all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); @@ -828,6 +850,7 @@ impl SegmentHolder { fetch_k, ef_search: graph_ef, filter: filter_arc.clone(), + tuning, reply: tx.clone(), }; if pool.submit(job) { @@ -934,15 +957,22 @@ impl SegmentHolder { for imm in &segments.immutable { let ef_seg = seg_ef(imm.suggested_ef()); if graph_filter.is_some() { - all.extend(imm.search_filtered( + all.extend(imm.search_filtered_with_tuning( query_f32, fetch_k, ef_seg, &mut snap.scratch, graph_filter, + tuning, )); } else { - let results = imm.search(query_f32, fetch_k, ef_seg, &mut snap.scratch); + let results = imm.search_with_tuning( + query_f32, + fetch_k, + ef_seg, + &mut snap.scratch, + tuning, + ); if post_filter { if let Some(bm) = filter_ref { all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); @@ -1242,6 +1272,7 @@ mod tests { dirty_set: &[], dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), }; let mvcc = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_ctx); @@ -1276,6 +1307,7 @@ mod tests { dirty_set: &[], dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), }; let results = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_ctx); assert_eq!(results.len(), 1); @@ -1333,6 +1365,7 @@ mod tests { dirty_set: std::slice::from_ref(&dirty_entry), dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), }; let results = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_ctx); @@ -1372,6 +1405,7 @@ mod tests { dirty_set: &[], dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), }; let r1 = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_empty); @@ -1383,6 +1417,7 @@ mod tests { dirty_set: &[], dimension: dim as u32, ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), }; let r2 = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_empty2); diff --git a/src/vector/segment/immutable.rs b/src/vector/segment/immutable.rs index 7733e13f..83fb6b3f 100644 --- a/src/vector/segment/immutable.rs +++ b/src/vector/segment/immutable.rs @@ -27,6 +27,7 @@ use crate::vector::turbo_quant::inner_product::{prepare_query_prod, score_l2_pro use crate::vector::turbo_quant::sq8::{decode_sq8, sq8_params}; use crate::vector::turbo_quant::sub_centroid; use crate::vector::types::SearchResult; +use crate::vector::types::SearchTuning; use crate::vector::types::VectorId; /// Microseconds since UNIX epoch, read from the shard's thread-local cached @@ -450,18 +451,25 @@ impl ImmutableSegment { /// Candidates still carry per-segment internal ids (pre /// `remap_to_global_ids`); `graph.to_bfs` maps them into the sidecar. /// - /// Cost control: only the top `4·k` ADC-ranked candidates are re-scored - /// (candidates arrive nearest-first from the beam). The true top-k landing - /// outside a 4× ADC oversample is rare; re-scoring the full ef-wide beam - /// costs ~ef·dim f16 decodes per segment for negligible recall beyond that. - fn rerank_exact(&self, candidates: &mut SmallVec<[SearchResult; 32]>, query: &[f32], k: usize) { + /// Cost control: only the top `mult·k` ADC-ranked candidates are re-scored + /// (candidates arrive nearest-first from the beam; `mult` defaults to 4, + /// runtime-tunable via FT.CONFIG RERANK_MULT). The true top-k landing + /// outside a 4× ADC oversample is rare; deeper oversampling recovers the + /// tail at ~mult·k·dim f16 decodes per segment. + fn rerank_exact( + &self, + candidates: &mut SmallVec<[SearchResult; 32]>, + query: &[f32], + k: usize, + mult: u32, + ) { let Some(raw) = self.raw_f16() else { return; }; if candidates.is_empty() { return; } - let rerank_n = (4 * k.max(1)).min(candidates.len()); + let rerank_n = ((mult.max(1) as usize) * k.max(1)).min(candidates.len()); let dim = self.collection_meta.dimension as usize; let is_l2 = self.collection_meta.metric == crate::vector::types::DistanceMetric::L2; @@ -566,17 +574,53 @@ impl ImmutableSegment { k: usize, ef_search: usize, scratch: &mut SearchScratch, + ) -> SmallVec<[SearchResult; 32]> { + self.search_with_tuning(query, k, ef_search, scratch, SearchTuning::default()) + } + + /// [`Self::search`] with explicit recall/QPS knobs (FT.CONFIG + /// RERANK_MULT / EXACT_BEAM). The knob-free wrapper uses the defaults. + pub fn search_with_tuning( + &self, + query: &[f32], + k: usize, + ef_search: usize, + scratch: &mut SearchScratch, + tuning: SearchTuning, ) -> SmallVec<[SearchResult; 32]> { // WS3 idle-unload: record this search so the HOT->WARM idle-eligibility // check has an accurate recency signal (see `idle_secs`). self.touch_last_access(); + // EXACT_BEAM: navigate the graph with true f16 distances from the + // exact sidecar — the beam's candidate SELECTION becomes exact, so + // recall is graph-limited rather than quantization-limited. Requires + // the sidecar; segments without one keep the quantized ADC beam. + let exact_f16 = if tuning.exact_beam { + self.raw_f16() + } else { + None + }; // Use sub-centroid signs during beam (32-level LUT) when available. // This eliminates the separate rerank pass — beam itself is high-accuracy. // Note: passing ef_search for both k and ef_search is intentional. // HNSW returns up to `ef_search` candidates (no early truncation to k). // This preserves candidates for cross-segment merging in the caller, // which does the final top-k selection after merging all segments. - let mut candidates = if !self.sub_centroid_signs.is_empty() { + let mut candidates = if exact_f16.is_some() { + hnsw_search_filtered( + &self.graph, + self.vectors_tq.as_slice(), + query, + &self.collection_meta, + ef_search, + ef_search, + scratch, + None, + &self.sub_centroid_signs, + self.sub_sign_bytes_per_vec, + exact_f16, + ) + } else if !self.sub_centroid_signs.is_empty() { hnsw_search_subcent( &self.graph, self.vectors_tq.as_slice(), @@ -606,19 +650,22 @@ impl ImmutableSegment { cands }; // Filter deleted entries first so tombstones neither consume the - // exact-rerank 4·k budget nor leave stale ADC scores mixed into the + // exact-rerank mult·k budget nor leave stale ADC scores mixed into the // post-rerank ordering. candidates.retain(|c| { let bfs = self.graph.to_bfs(c.id.0); self.is_live_bfs(bfs) }); - // HQ-1: exact rerank of the top 4·k live beam candidates from the f16 - // sidecar — replaces quantized estimates with true metric distances - // before top-k truncation. No-op without a sidecar. Runs for SQ8 - // too: an A/B (20k×384d, 1seg+5seg) measured skipping it costs - // R@10 −0.010 clustered / −0.017 gaussian-5seg — real recall, and - // the SIMD sidecar kernels make the pass ~3% of a query. - self.rerank_exact(&mut candidates, query, k); + // HQ-1: exact rerank of the top mult·k live beam candidates from the + // f16 sidecar — replaces quantized estimates with true metric + // distances before top-k truncation. No-op without a sidecar. Runs + // for SQ8 too: an A/B (20k×384d, 1seg+5seg) measured skipping it + // costs R@10 −0.010 clustered / −0.017 gaussian-5seg — real recall, + // and the SIMD sidecar kernels make the pass ~3% of a query. + // Skipped under EXACT_BEAM: the beam already produced true distances. + if exact_f16.is_none() { + self.rerank_exact(&mut candidates, query, k, tuning.rerank_mult); + } candidates.truncate(k); self.remap_to_global_ids(&mut candidates); candidates @@ -633,8 +680,34 @@ impl ImmutableSegment { scratch: &mut SearchScratch, allow_bitmap: Option<&RoaringBitmap>, ) -> SmallVec<[SearchResult; 32]> { - // WS3 idle-unload: see comment in `search()` above. + self.search_filtered_with_tuning( + query, + k, + ef_search, + scratch, + allow_bitmap, + SearchTuning::default(), + ) + } + + /// [`Self::search_filtered`] with explicit recall/QPS knobs. + pub fn search_filtered_with_tuning( + &self, + query: &[f32], + k: usize, + ef_search: usize, + scratch: &mut SearchScratch, + allow_bitmap: Option<&RoaringBitmap>, + tuning: SearchTuning, + ) -> SmallVec<[SearchResult; 32]> { + // WS3 idle-unload: see comment in `search_with_tuning()` above. self.touch_last_access(); + // EXACT_BEAM (see `search_with_tuning`): exact f16 beam navigation. + let exact_f16 = if tuning.exact_beam { + self.raw_f16() + } else { + None + }; // Note: passing ef_search for both k and ef_search is intentional // (see comment in search() method above). let mut candidates = hnsw_search_filtered( @@ -648,12 +721,13 @@ impl ImmutableSegment { allow_bitmap, &self.sub_centroid_signs, self.sub_sign_bytes_per_vec, + exact_f16, ); // When sub-centroid signs are used in beam, no rerank needed. // Only rerank if beam used standard 16-level scoring (and no exact // sidecar — rerank_exact below supersedes the estimator). - if self.sub_centroid_signs.is_empty() && self.raw_f16.is_none() { + if exact_f16.is_none() && self.sub_centroid_signs.is_empty() && self.raw_f16.is_none() { self.rerank_with_prod(&mut candidates, query); } // Filter deleted entries first (see comment in search()). @@ -662,8 +736,11 @@ impl ImmutableSegment { self.is_live_bfs(bfs) }); // HQ-1: exact rerank of the live beam from the f16 sidecar (see the - // recall A/B note in `search`). - self.rerank_exact(&mut candidates, query, k); + // recall A/B note in `search`). Skipped under EXACT_BEAM — the beam + // already produced true distances. + if exact_f16.is_none() { + self.rerank_exact(&mut candidates, query, k, tuning.rerank_mult); + } candidates.truncate(k); self.remap_to_global_ids(&mut candidates); candidates @@ -1076,6 +1153,19 @@ impl ImmutableSegment { self.collection_meta.fwht_sign_flips.as_slice(), ); + // L2-on-TQ correction: ||a-q||^2 = (na-nq)^2 + (nq/na)*(na^2*d_sphere^2). + let l2_adjust = self.collection_meta.metric == crate::vector::types::DistanceMetric::L2; + let tq_fin = |v: f32, na: f32| -> f32 { + if !l2_adjust { + return v; + } + let diff = na - q_norm; + if na <= 0.0 { + return diff * diff; + } + diff * diff + (q_norm / na) * v + }; + // Brute-force scan with max-heap for top-K. // TQ codes are in BFS order — use graph.to_original(bfs_pos) for original ID. let tq_buf = self.vectors_tq.as_slice(); @@ -1091,7 +1181,7 @@ impl ImmutableSegment { // Map BFS position → original ID (same mapping HNSW search uses) let original_id = self.graph.to_original(bfs_pos as u32); - let dist = tq_l2_adc_scaled(&q_rotated, code, norm, centroids); + let dist = tq_fin(tq_l2_adc_scaled(&q_rotated, code, norm, centroids), norm); if heap.len() < k { heap.push((ordered_float::OrderedFloat(dist), original_id)); diff --git a/src/vector/segment/ivf.rs b/src/vector/segment/ivf.rs index cfbf5c3c..b133a3ea 100644 --- a/src/vector/segment/ivf.rs +++ b/src/vector/segment/ivf.rs @@ -9,7 +9,7 @@ use smallvec::SmallVec; use crate::vector::aligned_buffer::AlignedBuffer; use crate::vector::distance::fastscan; -use crate::vector::turbo_quant::codebook::CENTROIDS; +use crate::vector::turbo_quant::codebook::scaled_centroids; use crate::vector::turbo_quant::encoder::padded_dimension; use crate::vector::types::SearchResult; @@ -140,51 +140,6 @@ pub fn interleave_posting_list( } } -/// Maximum possible single-coordinate squared distance for LUT quantization. -/// -/// Conservative bound: the largest FWHT coordinate for a unit vector is bounded, -/// and the largest centroid is CENTROIDS[15]. We use a generous bound. -const MAX_SINGLE_COORD_DIST_SQ: f32 = 0.03; - -/// Scale factor for quantizing float distances to u8. -const LUT_SCALE: f32 = 240.0 / MAX_SINGLE_COORD_DIST_SQ; - -/// Quantize a single float squared distance to u8 [0, 255]. -#[inline] -fn quantize_dist_to_u8(dist_sq: f32) -> u8 { - let scaled = dist_sq * LUT_SCALE; - if scaled >= 255.0 { - 255 - } else if scaled <= 0.0 { - 0 - } else { - scaled as u8 - } -} - -/// Precompute u8 distance LUT from a rotated query vector. -/// -/// For each coordinate `coord` in `0..padded_dim`, produces 16 entries: -/// `lut_out[coord * 16 + k] = quantize_dist_to_u8((q_rotated[coord] - CENTROIDS[k])^2)` -/// -/// `lut_out` must have length >= `padded_dim * 16`. -/// -/// No allocations. Caller provides output buffer. -#[inline] -pub fn precompute_lut(q_rotated: &[f32], lut_out: &mut [u8]) { - let padded_dim = q_rotated.len(); - debug_assert!(lut_out.len() >= padded_dim * 16); - - for coord in 0..padded_dim { - let q_val = q_rotated[coord]; - let base = coord * 16; - for k in 0..16 { - let diff = q_val - CENTROIDS[k]; - lut_out[base + k] = quantize_dist_to_u8(diff * diff); - } - } -} - /// An IVF segment: cluster centroids + posting lists of quantized vectors. pub struct IvfSegment { /// Flat array of cluster centroids: n_clusters * dimension floats. @@ -201,6 +156,10 @@ pub struct IvfSegment { padded_dim: u32, /// FWHT sign flips used to rotate queries before LUT precomputation. sign_flips: AlignedBuffer, + /// Dimension-scaled 16-centroid codebook for LUT construction. + /// MUST match the codebook used to encode the stored TQ codes + /// (`scaled_centroids(padded_dim)`, codebook v2). + lut_centroids: [f32; 16], } impl IvfSegment { @@ -213,14 +172,16 @@ impl IvfSegment { dimension: u32, sign_flips: AlignedBuffer, ) -> Self { + let padded_dim = padded_dimension(dimension); Self { centroids, posting_lists, n_clusters, quantization, dimension, - padded_dim: padded_dimension(dimension), + padded_dim, sign_flips, + lut_centroids: scaled_centroids(padded_dim), } } @@ -288,8 +249,17 @@ impl IvfSegment { nprobe: usize, lut_buf: &mut [u8], ) -> SmallVec<[SearchResult; 32]> { - // Precompute u8 distance LUT from rotated query. - precompute_lut(q_rotated, lut_buf); + // Build the quantized u8 distance LUT from the rotated query using the + // same dimension-scaled codebook the stored codes were encoded with. + // + // TODO(metric-faithful L2): this scorer still ranks by unit-sphere + // distance scaled by norm² and lacks the l2_adjust/tq_fin correction + // that mutable.rs and hnsw/search.rs apply for raw-L2 queries. IVF + // segments are not yet built on any production path (holder.ivf is + // only populated by tests) — the correction MUST be added here before + // wiring IVF into compaction, or L2 recall collapses on + // high-magnitude datasets (see the gist-960 TQ-L2 fix). + let lut_params = fastscan::build_quantized_lut(q_rotated, &self.lut_centroids, lut_buf); let dim = self.dimension as usize; let pdim = self.padded_dim as usize; @@ -314,6 +284,7 @@ impl IvfSegment { fastscan::scan_posting_list( pl.codes.as_slice(), lut_buf, + lut_params, dim_half, &pl.ids, &pl.norms, @@ -726,42 +697,45 @@ mod tests { } #[test] - fn test_precompute_lut_known_query() { - // Query: all zeros -> distance to each centroid k = CENTROIDS[k]^2 - let padded_dim = 4; - let q = vec![0.0f32; padded_dim]; - let mut lut = vec![0u8; padded_dim * 16]; - precompute_lut(&q, &mut lut); - - // For each coord (all zero), LUT entry k = quantize(CENTROIDS[k]^2) - for coord in 0..padded_dim { - for k in 0..16 { - let expected_dist = CENTROIDS[k] * CENTROIDS[k]; - let expected_u8 = quantize_dist_to_u8(expected_dist); - assert_eq!( - lut[coord * 16 + k], - expected_u8, - "LUT mismatch at coord={coord}, k={k}: dist={expected_dist}" - ); - } - // Centroid 7 and 8 are near zero, should have smallest distances - #[allow(clippy::identity_op)] - { - assert!(lut[coord * 16 + 7] <= lut[coord * 16 + 0]); + fn test_build_quantized_lut_known_query() { + // Query: all zeros -> distance to each centroid k = centroids[k]^2. + // The per-coordinate minimum (centroids 7/8, nearest to zero) must map + // to LUT entry 0; farther centroids must be monotonically larger. + let padded_dim = 4u32; + let centroids = scaled_centroids(padded_dim); + let q = vec![0.0f32; padded_dim as usize]; + let mut lut = vec![0u8; padded_dim as usize * 16]; + let params = fastscan::build_quantized_lut(&q, ¢roids, &mut lut); + + assert!(params.scale > 0.0); + // bias = sum of per-coord minima = padded_dim * centroids[8]^2 (nearest). + let expected_bias = padded_dim as f32 * centroids[8] * centroids[8]; + assert!((params.bias - expected_bias).abs() < 1e-6); + + for coord in 0..padded_dim as usize { + // Nearest centroids (7/8) -> residual 0. + assert_eq!(lut[coord * 16 + 7], lut[coord * 16 + 8]); + assert_eq!(lut[coord * 16 + 8], 0, "nearest centroid must map to 0"); + // Monotonically increasing away from center. + for k in 8..15 { + assert!(lut[coord * 16 + k] <= lut[coord * 16 + k + 1]); } - assert!(lut[coord * 16 + 8] <= lut[coord * 16 + 15]); + // Farthest centroid must hit the top of the u8 range (scale is + // chosen as 255/max_range here since padded_dim is tiny). + assert_eq!(lut[coord * 16 + 15], 255); } } #[test] - fn test_precompute_lut_symmetry() { - // Query at zero: CENTROIDS are symmetric, so LUT[k] == LUT[15-k] - let padded_dim = 2; - let q = vec![0.0f32; padded_dim]; - let mut lut = vec![0u8; padded_dim * 16]; - precompute_lut(&q, &mut lut); - - for coord in 0..padded_dim { + fn test_build_quantized_lut_symmetry() { + // Query at zero: centroids are symmetric, so LUT[k] == LUT[15-k]. + let padded_dim = 2u32; + let centroids = scaled_centroids(padded_dim); + let q = vec![0.0f32; padded_dim as usize]; + let mut lut = vec![0u8; padded_dim as usize * 16]; + fastscan::build_quantized_lut(&q, ¢roids, &mut lut); + + for coord in 0..padded_dim as usize { for k in 0..16 { assert_eq!( lut[coord * 16 + k], @@ -827,18 +801,6 @@ mod tests { assert_eq!(seg.total_vectors(), 30); } - #[test] - fn test_quantize_dist_to_u8_range() { - // Zero distance -> 0 - assert_eq!(quantize_dist_to_u8(0.0), 0); - // Max distance -> 240 - assert_eq!(quantize_dist_to_u8(MAX_SINGLE_COORD_DIST_SQ), 240); - // Over max -> clamped to 255 - assert_eq!(quantize_dist_to_u8(1.0), 255); - // Negative -> 0 - assert_eq!(quantize_dist_to_u8(-0.1), 0); - } - // ----------------------------------------------------------------------- // k-means tests // ----------------------------------------------------------------------- diff --git a/src/vector/segment/mutable.rs b/src/vector/segment/mutable.rs index 5c3d87bd..caa6e9f5 100644 --- a/src/vector/segment/mutable.rs +++ b/src/vector/segment/mutable.rs @@ -12,7 +12,9 @@ use roaring::RoaringBitmap; use smallvec::SmallVec; use crate::vector::distance; +use crate::vector::distance::fastscan; use crate::vector::mvcc::visibility::is_visible; +use crate::vector::segment::ivf::BLOCK_SIZE as FASTSCAN_BLOCK; use crate::vector::turbo_quant::collection::{CollectionMetadata, QuantizationConfig}; use crate::vector::turbo_quant::encoder::{ encode_tq_mse_a2, encode_tq_mse_scaled, encode_tq_mse_scaled_with_signs, padded_dimension, @@ -90,6 +92,25 @@ pub struct BruteForceQuery { sq8_qi8: Vec, sq8_q_scale: f32, sq8_sum_qi8: i32, + /// FastScan pre-filter (TQ-ADC path only): quantized u8 LUT built once at + /// prepare (`padded_dim * 16` entries). Empty when the segment is too + /// small or the path doesn't apply — the chunk scan then runs the plain + /// per-candidate loop. + /// L2-on-TQ metric correction (see `tq_finish` below): raw query norm + /// and whether the transform applies (TQ-ADC path + L2 metric). TQ ranks + /// by sphere_dist*norm_a^2, which is metric-unfaithful for unnormalized + /// L2 data (gist-960 recall collapse). + l2_adjust: bool, + l2_q_norm: f32, + fs_lut: Vec, + /// 1/scale from `build_quantized_lut` (f32 distance reconstruction). + fs_inv_scale: f32, + /// Bias from `build_quantized_lut` (sum of per-coordinate minima). + fs_bias: f32, + /// Sound quantization error bound (`padded_dim / scale`, pre-norm²): + /// `|reconstructed - exact| <= fs_eps`, so `reconstructed - fs_eps` is a + /// true lower bound and the filter can never drop a real top-k candidate. + fs_eps: f32, /// Shared top-k accumulator across all chunks. heap: BinaryHeap, } @@ -149,6 +170,12 @@ pub struct FrozenSegment { struct MutableSegmentInner { /// TQ-encoded codes for HNSW TQ-ADC traversal. tq_codes: Vec, + /// FastScan shadow of the nibble-packed TQ codes in FAISS-interleaved + /// 32-vector blocks (`block[d * 32 + lane]`, zero-padded lanes). Written + /// at append time for scalar-codebook TQ4 collections (empty for SQ8/A2); + /// lets the MVCC brute-force scan batch 32 candidates per SIMD LUT pass. + /// Costs padded_dim/2 bytes per vector on top of `tq_codes`. + fs_blocks: Vec, /// QJL sign bits per vector — for TurboQuant_prod unbiased IP scoring. /// Zero-filled at insert time; recomputed from raw_f32 during freeze(). qjl_signs: Vec, @@ -230,6 +257,53 @@ fn encode_sq8_slot(vector_f32: &[f32], metric: DistanceMetric, dim: usize) -> (V (codes, raw_norm) } +/// Minimum entry count before the FastScan pre-filter engages. Below this, +/// the ~2-8µs per-query LUT build costs more than the plain scan saves. +const FASTSCAN_MIN_ENTRIES: usize = 64; + +/// Mirror one vector's nibble-packed TQ code into the FastScan shadow blocks. +/// +/// `internal_id` must be the next sequential id (append-only invariant): a new +/// zero-filled block is grown whenever a block boundary is crossed, so lanes +/// past the current count stay zero (FastScan scans full blocks; out-of-range +/// lanes are skipped by position, never by content). +fn push_fastscan_shadow(fs_blocks: &mut Vec, internal_id: usize, code: &[u8]) { + let dim_half = code.len(); + let block_bytes = dim_half * FASTSCAN_BLOCK; + if internal_id % FASTSCAN_BLOCK == 0 { + fs_blocks.resize(fs_blocks.len() + block_bytes, 0); + } + let block_base = (internal_id / FASTSCAN_BLOCK) * block_bytes; + let lane = internal_id % FASTSCAN_BLOCK; + for (d, &b) in code.iter().enumerate() { + fs_blocks[block_base + d * FASTSCAN_BLOCK + lane] = b; + } +} + +impl MutableSegmentInner { + /// Single shadow-maintenance gate shared by every TQ-code writer + /// (`append` + `append_transactional`): mirrors the nibble code into + /// `fs_blocks` for scalar-codebook TQ4 and accounts the extra bytes. + /// Keeping this in ONE place is what guarantees `fs_blocks` can never + /// desync from `tq_codes` when a new insertion path is added. + fn maintain_fastscan_shadow(&mut self, internal_id: u32, code: &[u8], is_a2: bool) { + if !is_a2 && code.len() * 2 == self.padded_dimension as usize { + push_fastscan_shadow(&mut self.fs_blocks, internal_id as usize, code); + self.byte_size += code.len(); + } + } + + /// O(1) structural check that the shadow covers exactly the blocks needed + /// for `entries.len()` vectors. Used as the FastScan gate: a mismatch + /// (e.g. a future writer that misses the shadow) falls back to the plain + /// scan instead of scanning stale/garbage lanes. + #[inline] + fn fastscan_shadow_consistent(&self, code_len: usize) -> bool { + let expected = self.entries.len().div_ceil(FASTSCAN_BLOCK) * code_len * FASTSCAN_BLOCK; + !self.fs_blocks.is_empty() && self.fs_blocks.len() == expected + } +} + /// Append-only flat buffer with TQ-ADC brute-force search. pub struct MutableSegment { inner: RwLock, @@ -254,6 +328,7 @@ impl MutableSegment { Self { inner: RwLock::new(MutableSegmentInner { tq_codes: Vec::new(), + fs_blocks: Vec::new(), qjl_signs: Vec::new(), residual_norms: Vec::new(), raw_f32: Vec::new(), @@ -364,6 +439,10 @@ impl MutableSegment { inner.tq_codes.extend_from_slice(&code.codes); inner.tq_codes.extend_from_slice(&code.norm.to_le_bytes()); + // FastScan shadow: scalar-codebook TQ4 only (A2 never scans via ADC). + let is_a2 = self.collection.quantization == QuantizationConfig::TurboQuant4A2; + inner.maintain_fastscan_shadow(internal_id, &code.codes, is_a2); + // Append sub-centroid signs (Light mode TQ4 only) if let Some(signs) = sub_signs { inner.sub_centroid_signs.extend_from_slice(&signs); @@ -548,6 +627,23 @@ impl MutableSegment { None }; + // L2-on-TQ metric correction: TQ/A2 codes are unit DIRECTIONS + a norm + // trailer; sphere_dist*norm_a^2 is metric-unfaithful for unnormalized + // L2 data. Reconstruct ||a-q||^2 = (na-nq)^2 + (nq/na)*(na^2*d_sphere^2). + let l2_q_norm: f32 = query_f32.iter().map(|x| x * x).sum::().sqrt(); + let l2_adjust = + (use_tq_adc || use_a2_decoded_l2) && self.collection.metric == DistanceMetric::L2; + let tq_fin = |v: f32, na: f32| -> f32 { + if !l2_adjust { + return v; + } + let diff = na - l2_q_norm; + if na <= 0.0 { + return diff * diff; + } + diff * diff + (l2_q_norm / na) * v + }; + for entry in &inner.entries { if entry.delete_lsn != 0 { continue; @@ -583,12 +679,15 @@ impl MutableSegment { let d = q_rotated[j] - decoded[j]; sum += d * d; } - sum * norm_sq + tq_fin(sum * norm_sq, entry.norm) } else { f32::MAX } } else if use_tq_adc { - tq_l2_adc_scaled(&q_rotated, tq_code, entry.norm, centroids) + tq_fin( + tq_l2_adc_scaled(&q_rotated, tq_code, entry.norm, centroids), + entry.norm, + ) } else if let Some(qs) = query_state { let qjl_bpv = inner.qjl_bytes_per_vec; let qjl_offset = id * qjl_bpv; @@ -741,14 +840,46 @@ impl MutableSegment { } else { (Vec::new(), 0.0, 0) }; + // FastScan pre-filter LUT (TQ-ADC path only): one per-query build, + // gated on segment size so tiny segments skip the fixed LUT cost. + // `use_tq_adc` already excludes SQ8 and A2 (the paths without a + // scalar 16-centroid codebook / interleaved shadow). + let (fs_lut, fs_inv_scale, fs_bias, fs_eps) = + if use_tq_adc && self.inner.read().entries.len() >= FASTSCAN_MIN_ENTRIES { + let centroids = self.collection.codebook_16(); + let mut lut = vec![0u8; prepared.len() * 16]; + let params = fastscan::build_quantized_lut(&prepared, centroids, &mut lut); + // Round-to-nearest quantization: each of the padded_dim LUT entries + // carries at most 0.5 quanta of error, so 0.5·padded_dim/scale is a + // sound accumulated bound (kernel saturation only *under*-estimates, + // which routes the candidate to the exact rescore anyway). The 1e-3 + // headroom covers f32 accumulation order differences. + let eps = 0.5 * prepared.len() as f32 / params.scale * 1.001; + (lut, 1.0 / params.scale, params.bias, eps) + } else { + (Vec::new(), 0.0, 0.0, 0.0) + }; + + let l2_adjust = use_tq_adc && self.collection.metric == DistanceMetric::L2; + let l2_q_norm: f32 = if l2_adjust { + query_f32.iter().map(|x| x * x).sum::().sqrt() + } else { + 0.0 + }; BruteForceQuery { prepared, use_tq_adc, + l2_adjust, + l2_q_norm, sq8_q_sum, sq8_q_sumsq, sq8_qi8, sq8_q_scale, sq8_sum_qi8, + fs_lut, + fs_inv_scale, + fs_bias, + fs_eps, heap: BinaryHeap::with_capacity(k + 1), } } @@ -849,6 +980,128 @@ impl MutableSegment { let use_tq_adc = q.use_tq_adc; let q_rotated = q.prepared.as_slice(); + + // FastScan-filtered scan: batch 32 candidates per SIMD LUT pass, then + // exact-rescore only candidates whose sound lower bound + // `(approx - eps) * norm²` beats the current heap worst. Every pushed + // distance is still the exact `tq_l2_adc_scaled` value, so results are + // identical to the plain loop below — the filter only skips candidates + // that provably cannot enter the top-k. + // The shadow-consistency gate is the safety net: if a future writer + // ever misses `maintain_fastscan_shadow`, the O(1) length check fails + // and the scan falls back to the plain loop instead of reading + // stale/garbage lanes. Note: a `[lo,hi)` chunk that splits a 32-block + // re-scans that block's SIMD pass on the next chunk (out-of-range + // lanes are skipped by position) — the default yield chunking is + // 32-aligned, so this costs at most one duplicate block per + // operator-tuned unaligned boundary. + let use_fastscan = use_tq_adc + && !q.fs_lut.is_empty() + && code_len * 2 == q.prepared.len() + && inner.fastscan_shadow_consistent(code_len); + if use_fastscan { + let scan = fastscan::fastscan_dispatch().scan_block; + let fs_lut = q.fs_lut.as_slice(); + let (inv_scale, bias, eps) = (q.fs_inv_scale, q.fs_bias, q.fs_eps); + // L2-on-TQ correction (see BruteForceQuery::l2_adjust): both the + // final distance and the pre-filter LOWER BOUND transform through + // ||a-q||^2 = (na-nq)^2 + (nq/na)*v. The map is affine in v with + // non-negative slope, so a sound lower bound stays a sound lower + // bound after the transform. + let (l2_adjust, l2_qn) = (q.l2_adjust, q.l2_q_norm); + let tq_fin = |v: f32, na: f32| -> f32 { + if !l2_adjust { + return v; + } + let diff = na - l2_qn; + if na <= 0.0 { + return diff * diff; + } + diff * diff + (l2_qn / na) * v + }; + let heap = &mut q.heap; + let block_bytes = code_len * FASTSCAN_BLOCK; + let mut block_dists = [0u16; 32]; + + let mut block = lo / FASTSCAN_BLOCK; + while block * FASTSCAN_BLOCK < hi { + let vec_start = block * FASTSCAN_BLOCK; + let base = block * block_bytes; + scan( + &inner.fs_blocks[base..base + block_bytes], + fs_lut, + code_len, + &mut block_dists, + ); + + let lane_lo = lo.saturating_sub(vec_start); + let lane_hi = (hi - vec_start).min(FASTSCAN_BLOCK); + for lane in lane_lo..lane_hi { + let entry = &inner.entries[vec_start + lane]; + if !is_visible( + entry.insert_lsn, + entry.delete_lsn, + entry.txn_id, + snapshot_lsn, + my_txn_id, + committed, + ) { + continue; + } + if let Some(bm) = allow_bitmap { + let gid = inner.global_id_base + entry.internal_id; + if !bm.contains(gid) { + continue; + } + } + let norm_sq = entry.norm * entry.norm; + if heap.len() >= k { + let lower = tq_fin( + (block_dists[lane] as f32 * inv_scale + bias - eps) * norm_sq, + entry.norm, + ); + if let Some(&DistF32(worst, _, _)) = heap.peek() { + if lower >= worst { + continue; + } + } + } + + let id = entry.internal_id as usize; + let tq_offset = id * bytes_per_code; + let tq_code = &inner.tq_codes[tq_offset..tq_offset + code_len]; + let dist = tq_fin( + tq_l2_adc_scaled(q_rotated, tq_code, entry.norm, centroids), + entry.norm, + ); + + let global_id = inner.global_id_base + entry.internal_id; + if heap.len() < k { + heap.push(DistF32(dist, global_id, entry.key_hash)); + } else if let Some(&DistF32(worst, _, _)) = heap.peek() { + if dist < worst { + heap.pop(); + heap.push(DistF32(dist, global_id, entry.key_hash)); + } + } + } + block += 1; + } + return; + } + + // L2-on-TQ correction for the plain (non-FastScan) chunk loop. + let (l2_adjust, l2_qn) = (q.l2_adjust, q.l2_q_norm); + let tq_fin = |v: f32, na: f32| -> f32 { + if !l2_adjust { + return v; + } + let diff = na - l2_qn; + if na <= 0.0 { + return diff * diff; + } + diff * diff + (l2_qn / na) * v + }; let heap = &mut q.heap; for entry in &inner.entries[lo..hi] { @@ -873,7 +1126,10 @@ impl MutableSegment { let tq_code = &inner.tq_codes[tq_offset..tq_offset + code_len]; let dist = if use_tq_adc { - tq_l2_adc_scaled(q_rotated, tq_code, entry.norm, centroids) + tq_fin( + tq_l2_adc_scaled(q_rotated, tq_code, entry.norm, centroids), + entry.norm, + ) } else { let qs = query_state.unwrap(); let qjl_bpv = inner.qjl_bytes_per_vec; @@ -977,6 +1233,10 @@ impl MutableSegment { inner.tq_codes.extend_from_slice(&code.codes); inner.tq_codes.extend_from_slice(&code.norm.to_le_bytes()); + // FastScan shadow: scalar-codebook TQ4 only (shared gate with append()). + let is_a2 = self.collection.quantization == QuantizationConfig::TurboQuant4A2; + inner.maintain_fastscan_shadow(internal_id, &code.codes, is_a2); + let is_exact = self.collection.build_mode == crate::vector::turbo_quant::collection::BuildMode::Exact; crate::vector::f16::encode_f16_slice(vector_f32, &mut inner.raw_f16); @@ -1325,7 +1585,20 @@ impl MutableSegment { }) .collect(); + // ── FastScan shadow: rebuild (lanes rebase along with internal ids) ── + let fs_blocks = if inner.fs_blocks.is_empty() { + Vec::new() + } else { + let code_len = bpc - 4; + let mut fs = Vec::new(); + for i in 0..count { + push_fastscan_shadow(&mut fs, i, &tq_codes[i * bpc..i * bpc + code_len]); + } + fs + }; + let byte_size = count * (bpc + std::mem::size_of::()) + + fs_blocks.len() + count * sub_bpv + (if !qjl_signs.is_empty() { count * qjl_bpv @@ -1350,6 +1623,7 @@ impl MutableSegment { let new_inner = MutableSegmentInner { tq_codes, + fs_blocks, qjl_signs, residual_norms, raw_f32, @@ -1584,6 +1858,223 @@ mod tests { q_rot } + #[test] + fn test_fastscan_shadow_matches_vector_major_codes() { + distance::init(); + let dim = 32u32; + let col = make_collection(dim); + let seg = MutableSegment::new(dim, col); + + // 70 vectors: 2 full blocks + 1 partial (6 lanes). + for i in 0..70u32 { + let v = make_f32_vector(dim as usize, 1000 + i); + seg.append(i as u64, &v, i as u64 + 1); + } + + let inner = seg.inner.read(); + let code_len = inner.bytes_per_code - 4; + assert_eq!(code_len * 2, inner.padded_dimension as usize); + let block_bytes = code_len * FASTSCAN_BLOCK; + assert_eq!( + inner.fs_blocks.len(), + 3 * block_bytes, + "3 blocks for 70 entries" + ); + + for v in 0..70usize { + let code = + &inner.tq_codes[v * inner.bytes_per_code..v * inner.bytes_per_code + code_len]; + let block_base = (v / FASTSCAN_BLOCK) * block_bytes; + let lane = v % FASTSCAN_BLOCK; + for (d, &byte) in code.iter().enumerate() { + assert_eq!( + inner.fs_blocks[block_base + d * FASTSCAN_BLOCK + lane], + byte, + "shadow mismatch at vector {v}, sub-dim {d}" + ); + } + } + // Padding lanes of the partial block must be zero. + let last_base = 2 * block_bytes; + for d in 0..code_len { + for lane in 6..FASTSCAN_BLOCK { + assert_eq!( + inner.fs_blocks[last_base + d * FASTSCAN_BLOCK + lane], + 0, + "padding lane {lane} must stay zero" + ); + } + } + } + + #[test] + fn test_fastscan_filtered_scan_identical_to_plain_scan() { + distance::init(); + let dim = 32u32; + let col = make_collection(dim); + let seg = MutableSegment::new(dim, col); + + let n = 500u32; + for i in 0..n { + let v = make_f32_vector(dim as usize, 7000 + i); + seg.append(i as u64, &v, i as u64 + 1); + } + + let committed = roaring::RoaringTreemap::new(); + let query = make_f32_vector(dim as usize, 42); + let k = 10; + // Snapshot excludes the last 100 entries (insert_lsn > 400) and a + // bitmap filter drops every third id — exercises both skip paths + // inside the FastScan block loop. + let snapshot_lsn = 400u64; + let mut bitmap = RoaringBitmap::new(); + for i in 0..n { + if i % 3 != 0 { + bitmap.insert(i); + } + } + + // FastScan-filtered path (default: LUT built at prepare). + let mut q_fs = seg.prepare_brute_force_query(&query, false, k); + assert!( + !q_fs.fs_lut.is_empty(), + "FastScan LUT must engage above FASTSCAN_MIN_ENTRIES" + ); + seg.brute_force_scan_mvcc_chunk( + &mut q_fs, + None, + k, + Some(&bitmap), + snapshot_lsn, + 0, + &committed, + 0, + n as usize, + ); + let fs_results = q_fs.into_results(); + + // Plain path: identical query with the LUT cleared. + let mut q_plain = seg.prepare_brute_force_query(&query, false, k); + q_plain.fs_lut.clear(); + seg.brute_force_scan_mvcc_chunk( + &mut q_plain, + None, + k, + Some(&bitmap), + snapshot_lsn, + 0, + &committed, + 0, + n as usize, + ); + let plain_results = q_plain.into_results(); + + assert_eq!(fs_results.len(), plain_results.len()); + assert_eq!( + fs_results.len(), + k, + "enough visible entries for a full top-k" + ); + for (a, b) in fs_results.iter().zip(plain_results.iter()) { + assert_eq!(a.id, b.id, "FastScan filter changed the result set"); + assert_eq!(a.distance, b.distance, "distances must be bit-identical"); + } + + // Chunked scan (yield-path shape) must also match: same query fed in + // 64-entry chunks. + let mut q_chunked = seg.prepare_brute_force_query(&query, false, k); + let mut start = 0usize; + while start < n as usize { + let end = (start + 64).min(n as usize); + seg.brute_force_scan_mvcc_chunk( + &mut q_chunked, + None, + k, + Some(&bitmap), + snapshot_lsn, + 0, + &committed, + start, + end, + ); + start = end; + } + let chunked_results = q_chunked.into_results(); + for (a, b) in chunked_results.iter().zip(plain_results.iter()) { + assert_eq!(a.id, b.id, "chunked FastScan diverged"); + assert_eq!(a.distance, b.distance); + } + } + + #[test] + fn test_tq_l2_unnormalized_norm_variation_recall() { + // TQ encodes unit DIRECTIONS + a norm trailer. Ranking L2 by + // sphere_dist·‖a‖² is metric-unfaithful when norms vary (small-norm + // vectors look near regardless of direction) — the gist-960 recall + // collapse. The corrected score ‖a−q‖² = (‖a‖−‖q‖)² + ‖a‖‖q‖·d̂² + // must recover brute-force recall on unnormalized data. + distance::init(); + let dim = 32u32; + let col = make_collection(dim); // L2 + TQ4 + let seg = MutableSegment::new(dim, col); + + let n = 500u32; + let mut raw: Vec> = Vec::new(); + for i in 0..n { + // Unit direction scaled to norms spanning 0.5..40.5. + let mut v = make_f32_vector(dim as usize, 9000 + i); + let scale = 0.5 + (i % 41) as f32; + for x in v.iter_mut() { + *x *= scale; + } + seg.append(i as u64, &v, i as u64 + 1); + raw.push(v); + } + + let mut query = make_f32_vector(dim as usize, 777); + for x in query.iter_mut() { + *x *= 7.5; + } + + let k = 10; + let mut exact: Vec<(f32, u64)> = raw + .iter() + .enumerate() + .map(|(i, v)| { + let d: f32 = v + .iter() + .zip(query.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum(); + (d, i as u64) + }) + .collect(); + exact.sort_by(|a, b| a.0.total_cmp(&b.0)); + let want: std::collections::HashSet = exact[..k].iter().map(|(_, i)| *i).collect(); + + let committed = roaring::RoaringTreemap::new(); + let mut q = seg.prepare_brute_force_query(&query, false, k); + seg.brute_force_scan_mvcc_chunk( + &mut q, + None, + k, + None, + u64::MAX, + 0, + &committed, + 0, + n as usize, + ); + let results = q.into_results(); + let got: std::collections::HashSet = results.iter().map(|r| r.key_hash).collect(); + let overlap = want.intersection(&got).count(); + assert!( + overlap >= 7, + "TQ4+L2 recall on unnormalized data: {overlap}/{k} (norm-scaled \ + sphere ranking is metric-unfaithful)" + ); + } + #[test] fn test_append_returns_sequential_ids() { distance::init(); diff --git a/src/vector/store.rs b/src/vector/store.rs index b97562a3..ce6f2ed2 100644 --- a/src/vector/store.rs +++ b/src/vector/store.rs @@ -117,6 +117,56 @@ pub struct IndexMeta { /// field existed (v1-v3) default to `0` on load (legacy indexes become /// db-0-owned, matching pre-v0.6.0 global behavior for db 0 callers). pub db_index: u8, + /// Exact-rerank depth multiplier (HQ-1): re-score the top + /// `rerank_mult · k` beam candidates from the f16 sidecar. Default 4. + /// Range 1-64, set via FT.CONFIG SET RERANK_MULT. + pub rerank_mult: u32, + /// When true, the HNSW beam navigates with exact f16 sidecar distances + /// instead of quantized ADC (recall ≈ graph-limited, QPS cost grows with + /// dimension). Default false, set via FT.CONFIG SET EXACT_BEAM. + pub exact_beam: bool, +} + +/// Server-wide starting values for the per-index search-tuning knobs, +/// applied by FT.CREATE when the operator did not specify them +/// (`--vector-ef-runtime`, `--vector-rerank-mult`, `--vector-exact-beam`). +/// Per-index `FT.CONFIG SET` always overrides these at runtime. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct VectorCreateDefaults { + /// EF_RUNTIME for new indexes (0 = per-query auto heuristic). + pub ef_runtime: u32, + /// RERANK_MULT for new indexes (1-64). + pub rerank_mult: u32, + /// EXACT_BEAM for new indexes. + pub exact_beam: bool, +} + +impl Default for VectorCreateDefaults { + fn default() -> Self { + Self { + ef_runtime: 0, + rerank_mult: 4, + exact_beam: false, + } + } +} + +static VECTOR_CREATE_DEFAULTS: std::sync::OnceLock = + std::sync::OnceLock::new(); + +/// Install the server-wide FT.CREATE tuning defaults from `ServerConfig`. +/// Called once at startup (before any connection is accepted); a second call +/// is a no-op (the first write wins — matches `OnceLock` semantics, which is +/// what we want for embedded/test servers booting in one process). +pub fn set_vector_create_defaults(defaults: VectorCreateDefaults) { + let _ = VECTOR_CREATE_DEFAULTS.set(defaults); +} + +/// The server-wide FT.CREATE tuning defaults; the compiled-in baseline +/// (auto ef, mult 4, beam off) when the server never installed any — +/// unit tests and library embedders get pre-flag behavior unchanged. +pub fn vector_create_defaults() -> VectorCreateDefaults { + VECTOR_CREATE_DEFAULTS.get().copied().unwrap_or_default() } impl IndexMeta { @@ -2952,6 +3002,8 @@ pub(crate) fn test_index_meta(dim: u32) -> IndexMeta { merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } @@ -2981,6 +3033,8 @@ mod tests { merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } @@ -3003,6 +3057,8 @@ mod tests { merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } @@ -3698,6 +3754,8 @@ mod bg_compact_tests { merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } diff --git a/src/vector/turbo_quant/tq_adc.rs b/src/vector/turbo_quant/tq_adc.rs index 83c4de38..1a426d3a 100644 --- a/src/vector/turbo_quant/tq_adc.rs +++ b/src/vector/turbo_quant/tq_adc.rs @@ -537,6 +537,18 @@ pub fn tq_l2_adc_multibit_budgeted( if dist > budget { f32::MAX } else { dist } } +#[inline] +fn l2_true_from_sphere(sphere_v: f32, norm: f32, q_norm: f32) -> f32 { + // TQ codes are unit DIRECTIONS + a norm trailer; `sphere_v` = na^2 * d^2 + // on the unit sphere. That ranking is only metric-valid for unit-sphere + // metrics (COSINE/IP). True L2: ||a-q||^2 = (na-nq)^2 + (nq/na)*sphere_v. + let diff = norm - q_norm; + if norm <= 0.0 { + return diff * diff; + } + diff * diff + (q_norm / norm) * sphere_v +} + /// Brute-force scan of ALL TQ codes at any bit width using ADC. /// /// `bits`: quantization bit width (1-4). @@ -572,6 +584,7 @@ pub fn brute_force_tq_adc_multibit( *v *= inv; } } + let l2_adjust = collection.metric == crate::vector::types::DistanceMetric::L2; fwht::fwht( &mut q_rotated[..padded], collection.fwht_sign_flips.as_slice(), @@ -586,7 +599,12 @@ pub fn brute_force_tq_adc_multibit( let code = &tq_buffer[offset..offset + code_len]; let norm_bytes = &tq_buffer[offset + code_len..offset + code_len + 4]; let norm = f32::from_le_bytes([norm_bytes[0], norm_bytes[1], norm_bytes[2], norm_bytes[3]]); - let dist = tq_l2_adc_multibit(&q_rotated, code, norm, centroids, bits); + let sphere_v = tq_l2_adc_multibit(&q_rotated, code, norm, centroids, bits); + let dist = if l2_adjust { + l2_true_from_sphere(sphere_v, norm, q_norm) + } else { + sphere_v + }; if heap.len() < k { heap.push((ordered_float::OrderedFloat(dist), i as u32)); @@ -651,6 +669,7 @@ pub fn brute_force_tq_adc( *v *= inv; } } + let l2_adjust = collection.metric == crate::vector::types::DistanceMetric::L2; fwht::fwht( &mut q_rotated[..padded], collection.fwht_sign_flips.as_slice(), @@ -665,7 +684,12 @@ pub fn brute_force_tq_adc( let code = &tq_buffer[offset..offset + code_len]; let norm_bytes = &tq_buffer[offset + code_len..offset + code_len + 4]; let norm = f32::from_le_bytes([norm_bytes[0], norm_bytes[1], norm_bytes[2], norm_bytes[3]]); - let dist = tq_l2_adc_scaled(&q_rotated, code, norm, codebook); + let sphere_v = tq_l2_adc_scaled(&q_rotated, code, norm, codebook); + let dist = if l2_adjust { + l2_true_from_sphere(sphere_v, norm, q_norm) + } else { + sphere_v + }; if heap.len() < k { heap.push((ordered_float::OrderedFloat(dist), i as u32)); diff --git a/src/vector/types.rs b/src/vector/types.rs index f4b97113..6081e31a 100644 --- a/src/vector/types.rs +++ b/src/vector/types.rs @@ -52,6 +52,32 @@ impl SearchResult { } } +/// Per-query recall/QPS knobs threaded from `IndexMeta` (set via FT.CONFIG) +/// into the segment search paths. `Copy` so it rides in `MvccContext` / +/// `SearchSnapshot` at zero per-query cost. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SearchTuning { + /// Exact-rerank depth multiplier (HQ-1): the top `rerank_mult · k` + /// beam candidates are re-scored against the f16 sidecar. Default 4. + /// Higher recovers candidates the quantized ADC ranking dropped, at + /// ~mult·k·dim f16 decodes per segment. Range 1-64 (FT.CONFIG RERANK_MULT). + pub rerank_mult: u32, + /// When true and the segment has an f16 sidecar, the HNSW beam itself + /// navigates with exact f16 distances instead of quantized ADC — + /// recall becomes graph-limited (~Qdrant-parity) at a QPS cost that + /// grows with dimension. Default false (FT.CONFIG EXACT_BEAM ON|OFF). + pub exact_beam: bool, +} + +impl Default for SearchTuning { + fn default() -> Self { + Self { + rerank_mult: 4, + exact_beam: false, + } + } +} + impl Eq for SearchResult {} impl PartialOrd for SearchResult { diff --git a/tests/moonstore_warm_e2e.rs b/tests/moonstore_warm_e2e.rs index 5a509675..ebcf6958 100644 --- a/tests/moonstore_warm_e2e.rs +++ b/tests/moonstore_warm_e2e.rs @@ -38,6 +38,8 @@ fn make_test_meta(name: &str, dim: u32, compact_threshold: u32) -> IndexMeta { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } diff --git a/tests/mq_integration.rs b/tests/mq_integration.rs index 84b24ff0..83015b94 100644 --- a/tests/mq_integration.rs +++ b/tests/mq_integration.rs @@ -129,6 +129,11 @@ async fn start_mq_server(num_shards: usize) -> (u16, CancellationToken) { migrate_aof_from: None, migrate_aof_to: None, migrate_aof_shards: 0, + graph_result_cache_entries: 256, + graph_result_cache_bytes: 4_194_304, + vector_ef_runtime: 0, + vector_rerank_mult: 4, + vector_exact_beam: false, }; let cancel = token.clone(); diff --git a/tests/per_index_compaction_weight.rs b/tests/per_index_compaction_weight.rs index 4c33a626..236143e9 100644 --- a/tests/per_index_compaction_weight.rs +++ b/tests/per_index_compaction_weight.rs @@ -262,6 +262,8 @@ fn test_ft_config_weight_roundtrip() { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, }; store.create_index(meta).unwrap(); @@ -340,6 +342,8 @@ fn test_compaction_weight_persists_across_reload() { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, }; store.create_index(meta.clone()).unwrap(); @@ -393,6 +397,8 @@ fn test_vacuum_vector_weight_sets_and_reads() { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, }; store.create_index(meta).unwrap(); diff --git a/tests/quickwins_red.rs b/tests/quickwins_red.rs index 8164da55..bf3dc86a 100644 --- a/tests/quickwins_red.rs +++ b/tests/quickwins_red.rs @@ -48,6 +48,8 @@ mod qw7 { merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } diff --git a/tests/txn_kv_wiring.rs b/tests/txn_kv_wiring.rs index 76514571..4415c1da 100644 --- a/tests/txn_kv_wiring.rs +++ b/tests/txn_kv_wiring.rs @@ -134,6 +134,11 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can migrate_aof_from: None, migrate_aof_to: None, migrate_aof_shards: 0, + graph_result_cache_entries: 256, + graph_result_cache_bytes: 4_194_304, + vector_ef_runtime: 0, + vector_rerank_mult: 4, + vector_exact_beam: false, }; let cancel = token.clone(); diff --git a/tests/vector_config_defaults.rs b/tests/vector_config_defaults.rs new file mode 100644 index 00000000..4cc02aae --- /dev/null +++ b/tests/vector_config_defaults.rs @@ -0,0 +1,129 @@ +//! Server-wide FT.CREATE tuning defaults (`--vector-ef-runtime`, +//! `--vector-rerank-mult`, `--vector-exact-beam`) and the graph result-cache +//! limits (`--graph-result-cache-*`). +//! +//! Lives in its own integration binary ON PURPOSE: the defaults are +//! process-wide `OnceLock`s installed once at startup, so exercising the +//! installed path in the shared lib-test binary would leak the values into +//! every other vector test. Here the process is ours alone. + +use bytes::Bytes; + +use moon::command::vector_search::{ft_config, ft_create}; +use moon::protocol::Frame; +use moon::vector::store::{VectorCreateDefaults, VectorStore, set_vector_create_defaults}; + +fn bulk(b: &[u8]) -> Frame { + Frame::BulkString(Bytes::copy_from_slice(b)) +} + +fn ft_create_args(name: &[u8]) -> Vec { + vec![ + bulk(name), + bulk(b"ON"), + bulk(b"HASH"), + bulk(b"PREFIX"), + bulk(b"1"), + bulk(b"doc:"), + bulk(b"SCHEMA"), + bulk(b"vec"), + bulk(b"VECTOR"), + bulk(b"HNSW"), + bulk(b"6"), + bulk(b"TYPE"), + bulk(b"FLOAT32"), + bulk(b"DIM"), + bulk(b"8"), + bulk(b"DISTANCE_METRIC"), + bulk(b"L2"), + ] +} + +#[test] +fn server_defaults_apply_to_new_indexes_and_ft_config_overrides() { + // Simulates the startup install (main.rs / run_embedded). + set_vector_create_defaults(VectorCreateDefaults { + ef_runtime: 512, + rerank_mult: 16, + exact_beam: true, + }); + + let mut store = VectorStore::new(); + let mut text = moon::text::store::TextStore::new(); + let result = ft_create(&mut store, &mut text, &ft_create_args(b"defidx"), 0); + assert!(matches!(result, Frame::SimpleString(_)), "{result:?}"); + + #[allow(clippy::unwrap_used)] // index just created + let idx = store.get_index_mut_for_db(b"defidx", 0).unwrap(); + assert_eq!(idx.meta.hnsw_ef_runtime, 512, "--vector-ef-runtime default"); + assert_eq!(idx.meta.rerank_mult, 16, "--vector-rerank-mult default"); + assert!(idx.meta.exact_beam, "--vector-exact-beam default"); + + // FT.CONFIG GET reads the server-installed starting values. + let get = |store: &mut VectorStore, text: &mut moon::text::store::TextStore, param: &[u8]| { + ft_config( + store, + text, + &[bulk(b"GET"), bulk(b"defidx"), bulk(param)], + 0, + ) + }; + match get(&mut store, &mut text, b"RERANK_MULT") { + Frame::BulkString(b) => assert_eq!(&b[..], b"16"), + other => panic!("expected 16, got {other:?}"), + } + match get(&mut store, &mut text, b"EXACT_BEAM") { + Frame::BulkString(b) => assert_eq!(&b[..], b"ON"), + other => panic!("expected ON, got {other:?}"), + } + + // Per-index FT.CONFIG SET still overrides the server default. + let result = ft_config( + &mut store, + &mut text, + &[ + bulk(b"SET"), + bulk(b"defidx"), + bulk(b"EXACT_BEAM"), + bulk(b"OFF"), + ], + 0, + ); + assert!(matches!(result, Frame::SimpleString(_)), "{result:?}"); + #[allow(clippy::unwrap_used)] // index exists + let idx = store.get_index_mut_for_db(b"defidx", 0).unwrap(); + assert!(!idx.meta.exact_beam); + + // Explicit EF_RUNTIME in FT.CREATE wins over the server default. + let mut args = ft_create_args(b"efidx"); + // Grow the param count (6 -> 8) and append EF_RUNTIME 64. + let hnsw_cnt_pos = args + .iter() + .position(|f| matches!(f, Frame::BulkString(b) if &b[..] == b"6")) + .expect("param count"); + args[hnsw_cnt_pos] = bulk(b"8"); + args.push(bulk(b"EF_RUNTIME")); + args.push(bulk(b"64")); + let result = ft_create(&mut store, &mut text, &args, 0); + assert!(matches!(result, Frame::SimpleString(_)), "{result:?}"); + #[allow(clippy::unwrap_used)] // index just created + let idx = store.get_index_mut_for_db(b"efidx", 0).unwrap(); + assert_eq!(idx.meta.hnsw_ef_runtime, 64, "explicit EF_RUNTIME must win"); +} + +#[cfg(feature = "graph")] +#[test] +fn graph_result_cache_limits_are_configurable() { + use moon::graph::cypher::result_cache; + + // Before any install: compiled-in defaults. + // (This runs in the same process as the vector test above, but the two + // OnceLocks are independent; only this test touches the graph one. Test + // ordering within the binary does not matter for that reason.) + result_cache::set_configured_limits(64, 1 << 20); + assert_eq!(result_cache::configured_limits(), (64, 1 << 20)); + + // Second install is a no-op (first write wins). + result_cache::set_configured_limits(9999, 9999); + assert_eq!(result_cache::configured_limits(), (64, 1 << 20)); +} diff --git a/tests/vector_edge_cases.rs b/tests/vector_edge_cases.rs index b274fa0e..a9c2e7b0 100644 --- a/tests/vector_edge_cases.rs +++ b/tests/vector_edge_cases.rs @@ -44,6 +44,8 @@ fn make_meta(name: &str, dim: u32) -> IndexMeta { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } diff --git a/tests/vector_exact_rerank.rs b/tests/vector_exact_rerank.rs index 6ec8ea48..847cc6c3 100644 --- a/tests/vector_exact_rerank.rs +++ b/tests/vector_exact_rerank.rs @@ -58,6 +58,8 @@ fn make_meta(name: &str, quant: QuantizationConfig, metric: DistanceMetric) -> I merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } @@ -447,3 +449,224 @@ fn tombstoned_candidates_do_not_eat_rerank_budget() { let got: Vec = results.iter().map(|r| r.id.0 as usize).collect(); assert_eq!(got, expect, "live top-k mismatch after tombstone filtering"); } + +// ── 9. FT.CONFIG knobs: RERANK_MULT deepens the oversample ─────────────────── + +/// Search one compacted segment with explicit tuning knobs. +fn segment_search_tuned( + store: &mut VectorStore, + name: &str, + q: &[f32], + k: usize, + ef: usize, + tuning: moon::vector::types::SearchTuning, +) -> Vec { + let idx = store.get_index_mut(name.as_bytes()).expect("index"); + let snap = idx.segments.load(); + let seg = &snap.immutable[0]; + let mut scratch = moon::vector::hnsw::search::SearchScratch::new( + seg.total_count(), + padded_dimension(DIM as u32), + ); + seg.search_with_tuning(q, k, ef, &mut scratch, tuning) + .into_iter() + .collect() +} + +fn recall_at_k(results: &[SearchResult], truth_ids: &std::collections::HashSet) -> usize { + results + .iter() + .filter(|r| truth_ids.contains(&r.id.0)) + .count() +} + +fn true_top_k(q: &[f32], vecs: &[Vec], k: usize) -> std::collections::HashSet { + let mut truth: Vec<(f32, u32)> = vecs + .iter() + .enumerate() + .map(|(i, v)| (l2_sq(q, v), i as u32)) + .collect(); + truth.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + truth[..k].iter().map(|(_, i)| *i).collect() +} + +#[test] +fn rerank_mult_deeper_oversample_never_hurts_recall() { + use moon::vector::types::SearchTuning; + // TQ4's coarse codes misrank enough beam candidates that the exact + // rerank's oversample depth matters: with mult=1 only the top k ADC + // candidates are re-scored (true neighbors ranked k..4k by ADC are + // lost); deeper mult recovers them from the same beam. + let (mut store, vecs) = build_compacted( + "xr_mult", + QuantizationConfig::TurboQuant4, + DistanceMetric::L2, + 400, + 0x4EED5, + ); + let mut rng = Rng::new(0x517E); + let k = 10; + let (mut hits_m1, mut hits_m4, mut hits_m16, mut total) = (0usize, 0usize, 0usize, 0usize); + for _ in 0..20 { + let q = rng.vec(DIM); + let truth_ids = true_top_k(&q, &vecs, k); + for (mult, hits) in [ + (1u32, &mut hits_m1), + (4u32, &mut hits_m4), + (16u32, &mut hits_m16), + ] { + let tuning = SearchTuning { + rerank_mult: mult, + exact_beam: false, + }; + let results = segment_search_tuned(&mut store, "xr_mult", &q, k, 128, tuning); + *hits += recall_at_k(&results, &truth_ids); + } + total += k; + } + // Monotone: deeper oversample can only help (same beam, more re-scored). + assert!( + hits_m4 >= hits_m1, + "mult=4 recall {hits_m4}/{total} < mult=1 recall {hits_m1}/{total}" + ); + assert!( + hits_m16 >= hits_m4, + "mult=16 recall {hits_m16}/{total} < mult=4 recall {hits_m4}/{total}" + ); + // And the knob must actually bite on this seeded workload: shallow + // mult=1 measurably trails the mult=16 deep rerank. + assert!( + hits_m16 > hits_m1, + "mult=16 ({hits_m16}/{total}) did not improve over mult=1 \ + ({hits_m1}/{total}) — RERANK_MULT knob is not reaching rerank_exact" + ); +} + +// ── 10. FT.CONFIG knobs: EXACT_BEAM navigates with true f16 distances ──────── + +#[test] +fn exact_beam_recall_at_least_matches_adc_beam() { + use moon::vector::types::SearchTuning; + // A deliberately narrow beam (ef just above k) makes beam-candidate + // SELECTION the recall bottleneck: TQ4's quantized ADC steers the beam + // into wrong graph regions that no post-hoc rerank can repair. The + // exact f16 beam navigates with true distances — recall becomes + // graph-limited. Widely varying norms (0.5..40) are TQ4+L2's weak spot: + // the norm-corrected sphere estimator is unbiased but noisy there. + distance::init(); + let mut store = VectorStore::new(); + store + .create_index(make_meta( + "xr_beam", + QuantizationConfig::TurboQuant4, + DistanceMetric::L2, + )) + .expect("create_index"); + let mut vrng = Rng::new(0xBEA31); + let mut vecs: Vec> = Vec::with_capacity(400); + for i in 0..400usize { + let mut v = vrng.vec(DIM); + let scale = 0.5 + 39.5 * (vrng.next_f32() * 0.5 + 0.5); + for x in v.iter_mut() { + *x *= scale; + } + let key = Bytes::from(format!("doc:{i}")); + let key_hash = xxhash_rust::xxh64::xxh64(&key, 0); + store + .insert_vector(b"xr_beam", &v, key_hash, key) + .expect("insert"); + vecs.push(v); + } + store.force_compact_index(b"xr_beam").expect("compact"); + assert_eq!(store.immutable_segment_count(b"xr_beam"), Some(1)); + let mut rng = Rng::new(0x90D0); + let k = 10; + let ef = 16; + let (mut hits_adc, mut hits_exact, mut total) = (0usize, 0usize, 0usize); + for _ in 0..20 { + let q = rng.vec(DIM); + let truth_ids = true_top_k(&q, &vecs, k); + let adc = segment_search_tuned(&mut store, "xr_beam", &q, k, ef, SearchTuning::default()); + let exact = segment_search_tuned( + &mut store, + "xr_beam", + &q, + k, + ef, + SearchTuning { + rerank_mult: 4, + exact_beam: true, + }, + ); + hits_adc += recall_at_k(&adc, &truth_ids); + hits_exact += recall_at_k(&exact, &truth_ids); + total += k; + + // EXACT_BEAM distances are true metric values (f16 tolerance) — + // the beam scored them exactly, no rerank pass needed. + for r in &exact { + let truth = l2_sq(&q, &vecs[r.id.0 as usize]); + let rel = (r.distance - truth).abs() / truth.max(1e-6); + assert!( + rel < 1.5e-3, + "EXACT_BEAM id={} returned={} exact={truth} rel={rel} — beam \ + returned a quantized estimate", + r.id.0, + r.distance + ); + } + } + assert!( + hits_exact >= hits_adc, + "exact beam recall {hits_exact}/{total} < ADC beam recall {hits_adc}/{total}" + ); + // The knob must bite: at ef=16 the TQ4 ADC beam measurably trails. + assert!( + hits_exact > hits_adc, + "exact beam ({hits_exact}/{total}) did not improve over the ADC beam \ + ({hits_adc}/{total}) — EXACT_BEAM is not reaching the distance closures" + ); +} + +#[test] +fn exact_beam_without_sidecar_falls_back_to_adc() { + use moon::vector::persistence::segment_io::{read_immutable_segment, write_immutable_segment}; + use moon::vector::types::SearchTuning; + + let (mut store, _vecs) = build_compacted( + "xr_bleg", + QuantizationConfig::Sq8, + DistanceMetric::L2, + 32, + 0x1B3A, + ); + let dir = std::env::temp_dir().join(format!("moon-xr-bleg-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + { + let idx = store.get_index_mut(b"xr_bleg").expect("index"); + let snap = idx.segments.load(); + write_immutable_segment(&dir, 5, &snap.immutable[0], &idx.collection).expect("write"); + } + let _ = std::fs::remove_file(dir.join("segment-5").join("raw_f16.bin")); + + let (seg, _meta) = read_immutable_segment(&dir, 5).expect("read legacy segment"); + let mut scratch = moon::vector::hnsw::search::SearchScratch::new(seg.total_count(), 64); + let mut rng = Rng::new(0xFA11); + let q = rng.vec(DIM); + let results = seg.search_with_tuning( + &q, + 5, + 64, + &mut scratch, + SearchTuning { + rerank_mult: 4, + exact_beam: true, + }, + ); + assert!( + !results.is_empty(), + "EXACT_BEAM on a sidecar-less segment must silently fall back to ADC" + ); + let _ = std::fs::remove_dir_all(&dir); +} diff --git a/tests/vector_insert_bench.rs b/tests/vector_insert_bench.rs index 7fd567ff..fd489383 100644 --- a/tests/vector_insert_bench.rs +++ b/tests/vector_insert_bench.rs @@ -147,6 +147,8 @@ fn bench_full_insert_pipeline_128d() { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, }; let _ = store.create_index(meta); @@ -227,6 +229,8 @@ fn bench_full_insert_pipeline_768d() { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, }; let _ = store.create_index(meta); diff --git a/tests/vector_segment_merge.rs b/tests/vector_segment_merge.rs index 06ca7d4d..16539c5d 100644 --- a/tests/vector_segment_merge.rs +++ b/tests/vector_segment_merge.rs @@ -80,6 +80,8 @@ fn make_meta(name: &str, dim: u32) -> IndexMeta { merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } @@ -104,6 +106,8 @@ fn make_meta_sq8(name: &str, dim: u32) -> IndexMeta { merge_mode: MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } diff --git a/tests/vector_stress.rs b/tests/vector_stress.rs index 61881f48..8fbd9596 100644 --- a/tests/vector_stress.rs +++ b/tests/vector_stress.rs @@ -60,6 +60,8 @@ fn make_index_meta(name: &str, dim: u32) -> IndexMeta { merge_mode: moon::vector::segment::compaction::MergeMode::GraphUnion, keep_raw: false, db_index: 0, + rerank_mult: 4, + exact_beam: false, } } diff --git a/tests/workspace_integration.rs b/tests/workspace_integration.rs index a2cd6619..9d7df377 100644 --- a/tests/workspace_integration.rs +++ b/tests/workspace_integration.rs @@ -122,6 +122,11 @@ async fn start_workspace_server(num_shards: usize) -> (u16, CancellationToken) { migrate_aof_from: None, migrate_aof_to: None, migrate_aof_shards: 0, + graph_result_cache_entries: 256, + graph_result_cache_bytes: 4_194_304, + vector_ef_runtime: 0, + vector_rerank_mult: 4, + vector_exact_beam: false, }; let cancel = token.clone(); @@ -357,6 +362,11 @@ async fn start_workspace_server_with_auth( migrate_aof_from: None, migrate_aof_to: None, migrate_aof_shards: 0, + graph_result_cache_entries: 256, + graph_result_cache_bytes: 4_194_304, + vector_ef_runtime: 0, + vector_rerank_mult: 4, + vector_exact_beam: false, }; let cancel = token.clone();