diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/README.md b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/README.md new file mode 100644 index 0000000000..df3ffe763d --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/README.md @@ -0,0 +1,327 @@ +# Billion-Scale Synthetic Data Generator + +A synthetic vector-dataset generator for approximate-nearest-neighbor (ANN) benchmarking, built so that an index (Vamana / DiskANN, CAGRA, …) **behaves the same on the synthetic data as on the real data it was fit from** — matching both the **search recall–QPS Pareto** *and* the **index build time**. + +## Contents + +- **[Background](#background)** — The paper and the shipped baseline this improves on. +- **[1. Core idea](#1-core-idea)** — The fit/generate workflow. +- **[2. How to run](#2-how-to-run)** — Requirements and how to run the pipeline. +- **[3. Headline results](#3-headline-result)** — Some benchmark numbers. +- **[4. Parameters and their effects](#4-parameters-and-their-effects)** — How params affect the result. +- **[5. Block Decode mode](#5-block-decode-mode)** — streaming generate+decode that removes the default path's memory walls. +- **[6. TODOs and future work](#6-todos-and-future-work)** — cheap GT, broader validation, privacy. + +--- + +## Background + +This repo assumes you already understand the motivation for our billion-scale synthetic data generator and our existing methodology shipped in `cuvs_bench`. + +- **The Paper [(link)](Mimicking_Vector_Datasets_at_Billion_Scale.pdf)** — motivation, and the shipped generator's design and evaluation. +- **`cuvs_bench.synthesize-dataset` [(README)](https://github.com/NVIDIA/cuvs/blob/main/fern/pages/cuvs_bench/synthesize_dataset.md) [(Code)](https://github.com/NVIDIA/cuvs/tree/main/python/cuvs_bench/cuvs_bench/synthesize_dataset)** — the implementation of the generator described in the paper. + +**The problem this repo exists to solve:** The shipped generator focused on matching the search recall–QPS Pareto but is unaware of index build time. This repo aims to generate synthetic data that matches **both the recall–QPS Pareto and the build time**. + +--- + +## 1. Core idea + +The shipped generator matches vector coordinates. This repo instead matches the **navigability of the kNN graph** — which is what actually drives both build time and search. We model the kNN graph and then decode it back into vectors. + +### Workflow + +![Workflow](figures/workflow.png) + +We **fit** on a small `SS`-node real sample and **generate** a large `N`-node synthetic dataset: + +**Fit** — once, on the `SS`-node sample: +1. Build the sample's all-neighbors kNN graph and cluster it (KMeans). +2. Measure the graph statistics the upsampled kNN must preserve: cluster sizes, in-degree (hub) tail, in-cluster edge fraction, and the coherence (a.k.a. clustering coefficient — how often my neighbors' neighbors are also my neighbors, i.e. triangle density). +3. Train the decoder (a small neural net of MLPs) to map `(kNN graph + per-node features) → vector`, by regressing each sample node's real embedding from its graph neighborhood. +4. Fit the residual model — a per-cluster low-rank Gaussian (same as the shipped generator in `cuvs_bench`) capturing the within-cluster spread the decoder can't reproduce. + +**Generate** — produce `N` vectors in two stages: +1. **Upsampling stage** (`upsample.py`) — build a synthetic `N`-node kNN graph from the fitted statistics. +2. **Decoding stage** (`knn_decoder.py`) — run the trained decoder over that graph to produce `N` vectors (plus residual). + +The ANN index then builds and searches over the `N` synthetic vectors; the goal is that its build time and recall–QPS curve match the same index built on `N` *real* vectors. + +The two generate stages are the heart of the method, so the rest of this section explains each. + +### The Upsampling Stage: Building a coherent `N`-node graph + +- *Input:* the fitted statistics of the kNN graph built on the sample data. +- *Output:* an `N`-node kNN graph with the same stats. + +The naive approach is to wire the graph straight from the marginal statistics (cluster sizes, in-degree tail, in-cluster edge fraction). That reproduces build time but fails on recall, because drawing edges independently from marginals make a node's neighbors random picks within a cluster. However, in real kNNs, a point's neighbors are close to each other too. Without modeling this local coherence the search recall caps far below real. + +So we build coherence straight into the edges. Inside each cluster, lay the nodes in a line and wire each one to a few others in a small sliding window. Because neighboring windows overlap, neighbors end up sharing neighbors, naturally forming triangles. The window width is auto-tuned so the coherence matches the sample's. It's `O(N)` and needs no global kNN pass. Refer to `triadic_coherent_knn` function in `upsample.py`. + +On top of the coherent edges, a **[Chung–Lu hub](https://link.springer.com/article/10.1007/pl00012580)** backbone (weights drawn from the sample's measured in-degree distribution) is blended in, tuned by `--knn-frac`, to restore the heavy in-degree tail that plain coherent wiring misses. Refer to `blend_chunglu` function in `upsample.py`. + +Coherent edges supply the local structure, and the Chung–Lu blend supplies the hub tail and search difficulty. + +### The Decoding Stage: Turn the graph into vectors + +- *Input:* an `N`-node synthetic kNN graph (from the upsampling stage). +- *Output:* a `d`-dim vector per node. + +We use a **`kNNDecoder`**, which is a small neural net. Decoding runs in three steps: + +**1. Turn the kNN graph into per-node features.** The decoder never sees raw coordinates. It only sees the graph. Each synthetic node is assigned an anchor (which is one of the real sample points), and we build its feature vector from two things: + - the anchor embedding — `MLP(anchor)` of that sample point (roughly, which region of space the node lives in), and + - structural features read off the graph — the node's in-degree and its neighbors' mean in-degree (how *hub-like* the node and its neighborhood are). + +**2. Decode with the model.** For each node in the target `N` nodes, the `kNNDecoder` attention-pools its `k` neighbors' features, concatenates with the node's own features, and passes them through the MLP to output a `d`-dim vector for that node. It was trained (during Fit) with MSE to regress each node's real embedding. But many nodes map to nearly identical inputs (same cluster + similar local graph topology) yet had different real embeddings, so the decoder can only predict the conditional mean for that input (the average of all real vectors whose nodes shared that graph-position). The decoded manifold thus comes out very smooth, dropping the spread among nodes that share the same graph-position. + +**3. Add the residual → the spread.** To overcome the smoothness, for each node, draw a sample from a fitted low-rank Gaussian, like we do in the shipped `cuvs_bench` data synthesizer, and add it to the decoded mean. Finally, a radial percentile step resets each vector's norm from its cluster's real norm inverse-CDF (also as we already do in `cuvs_bench`). + +--- + +## 2. How to run + +### Requirements +An NVIDIA GPU, plus: +- **RAPIDS cuVS ≥ 26.06** — `cuvs.cluster.kmeans`, `cuvs.neighbors.all_neighbors` / `nn_descent` / `vamana`. +- **PyTorch**, **CuPy**, **NumPy**. (Easiest: a RAPIDS conda env with PyTorch installed into it.) +- (To reproduce results below): DiskANN. Build [microsoft/DiskANN](https://github.com/microsoft/DiskANN). + + ```bash + git clone https://github.com/microsoft/DiskANN && cd DiskANN + mkdir build && cd build && cmake .. && make -j + export DISKANN_APPS=$PWD/apps + ``` + +### The pipeline +**Generate data → build the index (build time) → search (recall/QPS)** — run once for synthetic, once for a real reference, then compare. + +**1. Real reference.** Omit `--target` to split the sample itself into base + held-out queries + exact GT: +```bash +python generate_data.py --sample SAMPLE.fbin --out-dir out/real_10m +``` + +**2. Synthetic dataset.** `--target` = number of base vectors; `--n-queries` (default 10000) held-out queries are generated on top (pool = target + n_queries). : +```bash +# 50K sample -> 10M +python generate_data.py --sample SAMPLE_50K.fbin --out-dir out/synth_10m \ + --target 9990000 --nc 500 --resid-rank 32 \ + --resid-scale 1.8 --knn-frac 0.5 --model-cache out/model_cache +``` +Each run writes `base.fbin`, `queries.fbin`, `groundtruth.neighbors.ibin`, `groundtruth.distances.fbin`. `--model-cache DIR` caches the trained decoder so repeated runs on the same sample skip the retrain. `python generate_data.py --help` for the rest. + +**3. Build the index + measure build time** : +Using the `--save` flag writes `vamana.index` next to the `base.fbin`. +```bash +python build_vamana.py out/synth_10m/base.fbin --num-vectors 9990000 --save +python build_vamana.py out/real_10m/base.fbin --num-vectors 9990000 --save +``` + +For data too big to build on GPU, call DiskANN's on-disk builder directly over `base.fbin`: +```bash +$DISKANN_APPS/build_disk_index --data_type float --dist_fn l2 \ + --data_path out/synth_100m/base.fbin \ + --index_path_prefix out/synth_100m/disk_index \ + -R 64 -L 128 -B 8 -M 1000 -QD 192 +``` + +**4. Search / recall–QPS.** The built `vamana.index` loads directly into DiskANN's `search_memory_index`. Pass it as `--index_path_prefix` and sweep the search width `L`, scoring against the bundle's ground truth. Run on both the synthetic and real bundles and compare Recall / QPS across the sweep: +```bash +$DISKANN_APPS/search_memory_index --data_type float --dist_fn l2 \ + --index_path_prefix out/synth_10m/vamana.index \ + --query_file out/synth_10m/queries.fbin \ + --gt_file out/synth_10m/groundtruth.bin \ + --recall_at 10 -L 10 20 30 40 50 100 200 300 \ + --result_path out/synth_10m/res +``` + +Similar for a large index, but use `search_disk_index` instead. + +> DiskANN expects the ground truth as a single truthset file. The bundle writes it as two files (`groundtruth.neighbors.ibin` + `groundtruth.distances.fbin`), so merge them into DiskANN's `[npts, dim]` header + uint32 ids + float32 distances layout first. + +### What's in the repo +- **`generate_data.py`** — entry point: fit on a sample, orchestrate the upsampling → decoding stages, write the data. +- **`upsample.py`** — the **upsampling stage**: build the `N`-node graph from the sample graph. +- **`knn_decoder.py`** — the **decoding stage** (`kNNDecoder`): graph → vectors (mean + per-cluster residual + radial norm), plus training. +- **`utils.py`** — shared helpers. +- **`build_vamana.py`** — additional helper that builds a cuVS Vamana index over a `base.fbin`, reports build time. Is irrelevant to the synthetic data generation pipeline. + +--- + +## 3. Headline Result + +Both the **recall–QPS search curve** and the **index build time** track real at 200× scale. All experiments are with Falcon: + +![Synthetic vs real — recall–QPS curve and build-time match on Falcon](figures/results_comparison.png) + +The raw numbers behind the figure: + +**50K → 10M** — build time: real **1,040 s** vs synth **1,035 s** + +| Search width `L` | Real 10M — Recall@10 (QPS) | Synthetic 50K→10M — Recall@10 (QPS) | +|:---:|:---:|:---:| +| 10 | 82.88% (31,387) | 82.00% (29,904) | +| 20 | 90.69% (19,719) | 92.55% (19,570) | +| 30 | 93.63% (16,958) | 96.22% (15,564) | +| 40 | 95.31% (14,493) | 97.80% (13,519) | +| 50 | 96.35% (12,617) | 98.62% (11,966) | +| 100 | 98.22% (7,481) | 99.74% (7,326) | +| 200 | 99.24% (4,527) | 99.97% (4,386) | +| 300 | 99.50% (3,255) | 99.98% (3,247) | + +**500K → 100M** — build time: real **3,455 s** vs synth **3,562 s** + +| Search width `L` | Real 100M — Recall@10 (QPS) | Synthetic 500K→100M — Recall@10 (QPS) | +|:---:|:---:|:---:| +| 10 | 68.04% (24,410) | 61.33% (23,366) | +| 20 | 85.40% (14,100) | 84.53% (13,692) | +| 30 | 90.43% (11,218) | 91.20% (10,098) | +| 40 | 92.70% (10,089) | 94.36% (8,528) | +| 50 | 94.15% (7,795) | 96.14% (7,797) | +| 100 | 96.84% (4,315) | 99.08% (4,167) | +| 200 | 98.30% (2,335) | 99.79% (2,220) | +| 300 | 98.77% (1,647) | 99.89% (1,416) | + +To reproduce the datasets use the config below: +```bash +# 50K -> 10M +python generate_data.py --sample SAMPLE_50K.fbin --out-dir out/50k_10m \ + --target 9990000 --nc 500 --resid-rank 32 \ + --resid-scale 1.8 --knn-frac 0.5 --model-cache out/model_cache +python build_vamana.py "$d/base.fbin" --num-vectors 9990000 --save # Build using cuVS vamana +# Then trasfered to search on DiskANN + +# 500K -> 100M +python generate_data.py --sample SAMPLE_500K.fbin --out-dir out/500k_100m \ + --target 99990000 --nc 5000 --resid-rank 32 \ + --resid-scale 1.6 --knn-frac 0.5 --model-cache out/model_cache +# Built and searched using DiskANN +``` + +--- + +## 4. Parameters and their effects + +This section shares some important knobs and intuition. + +### (a) Sample size: the difficulty dial + +We recommend using a sample of size `target / 200` or larger. The results in [§3](#3-headline-result) was achieved at this ratio (50K→10M and 500K→100M are both 200×). + +### (b) `--knn-frac`: Coherent vs Chung–Lu edge mix +`--knn-frac` sets each node's mix of **coherent** (within-cluster) edges and **Chung–Lu hub** (random) edges. + +- **`1.0`** = all coherent: easiest — highest recall, fewest hubs, fastest build. +- **Lower** = more random hub edges: harder search (more distance-comparisons), a heavier hub tail, and higher build time — but recall drops as coherence dilutes. + +So `--knn-frac` trades recall against difficulty (and build time). A sweep at 100K→10M: + + | `knn-frac` | build | R@10 @ L=10 (# avg comps) | @ L=50 | @ L=100 | + |---|---|---|---|---| + | 0.8 (more coherent) | 878 s | 89.49 (1025) | 99.50 | 99.93 | + | **0.5** | **1028 s** | **81.57** (1140) | 98.61 | 99.70 | + | 0.2 (more random) | 1127 s | 63.61 (1251) | 93.75 | 98.42 | + | *real 10M* | *1040 s* | *82.88* (1010) | *96.35* | *98.22* | + +**Takeaway:** `knn_frac` is the coherence↔difficulty balance (moves recall, build, *and* dist-comps together). Push it **up** to make the data easier if recall is too low / build too slow; push it **down** to make the data harder if recall is too high / build too fast. + + +### (c) `--nc`: Residual/norm cluster count (cuvs-bench's `nc`) + +`--nc` sets the number of clusters for the per-cluster residual/norm fit — same as cuvs-bench's `nc`. We recommend setting `--nc ≈ ss/100` (`ss` = sample size), which is also the default if omitted. + + +### (d) `--resid-scale`: Roughness dial +`--resid-scale` multiplies the per-cluster residual added on top of the decoder's mean vector. It's the **roughness** dial for moving recall. + +- **`0`** = decoder mean only: too smooth → recall overshoots real. +- **Higher** = rougher manifold → recall drops and dist-comps rise (most visibly at low `L`), and build time creeps up. + +Sweep at 100K→10M (effect concentrated at low `L`): + + | `resid-scale` | build | R@10 @ L=10 (comps) | @ L=50 | @ L=100 | + |---|---|---|---|---| + | **1.7** | **1028 s** | **81.57** (1140) | 98.61 | 99.70 | + | 1.8 | 1057 s | 76.93 (1164) | 98.21 | 99.74 | + | 2.0 | 1095 s | 72.17 (1210) | 97.26 | 99.56 | + | *real 10M* | *1040 s* | *82.88* (1010) | *96.35* | *98.22* | + + +**Takeaway:** Push it **up** if the data is too easy (recall too high, too few comps), **down** if recall is too low. + + +### (e) `--resid-rank`: the rank of the per-cluster residual Gaussian +`--resid-rank` is the rank `r` of the per-cluster low-rank Gaussian (PPCA) — how many principal directions the within-cluster spread is modeled in. This is the same `ncomp` knob as the shipped `cuvs_bench` generator and the paper, and it behaves here exactly as the paper reports. + +- **Effect:** raising `r` moves the manifold's intrinsic dimension (LID), build time, and dist-comps toward real. However, recall collapses, because the data gets harder as its intrinsic dimension grows. + +**Takeaway:** Just leave `--resid-rank` at 32 and don't use it to tune unless you have done thorough verification. Reach for `resid-scale` / `nc` instead (same conclusion as the paper) for tuning. + + + +--- + +## 5. Block Decode mode + +The default pipeline builds the whole synthetic graph `(N, k)` and then decodes it in one pass. However, it holds the full kNN graph and additional `O(N)` states in memory, so it'll be challenging reach the 100B scale this generator targets. `--block-local` is an experimental streaming mode ([`block_local.py`](block_local.py)) that removes this wall. + +### Where the memory goes + +**Fit / train — fine at any target.** The model is trained once on the `SS`-node sample data, which is much smaller (suggested `SS = N/200`) than the target `N`. When the sample itself becomes too big for GPU memory, the **`--host-gather`** flag keeps the training anchor table / kNN graph / structural features on the CPU RAM and ships each minibatch to the GPU per step. This improves memory usage at the cost of performance, but is runnable. + +**Generate + decode — the wall.** The default whole-graph path materializes everything at once: +- **Host RAM** — the full `(N, k)` kNN graph plus the `O(N)` per-node feature arrays (each node's anchor id, structural features, and target norm). +- **GPU** — the anchor-embedding table `cluster_mlp(anchors)`, shape `(SS, d_emb)`, where `SS ≈ N/200` and default `d_emb = 64`. + +At `N=100B` neither fits — and because this is the part that grows with `N`, it's the wall Block Decode targets. + +### How Block Decode removes the wall + +**Main idea:** group the `N` nodes by their anchor (the sample point each synthetic node is assigned to), and process one contiguous block of anchors at a time (the reason we say a few blocks of anchors is to improve GPU utilization, because the points assigned to a sample point - i.e. the anchor - is on average `N/SS ≈ 200`). +Coherent edges are within-anchor (the windowed wiring from [§1](#the-upsampling-stage-building-a-coherent-n-node-graph)), so every node's coherent neighbors fall in the same block. Each block generates + decodes on its own, streams its base rows to disk, and is discarded. The full graph, the global sort, and the `O(N)` arrays never materialize; peak memory is set by the block size (`--block-size`), not `N`. + +**The remaining wall is the hub tail.** Coherent edges are local. They stay inside a node's anchor, hence inside its block that will be generated together with the approach above. However, the Chung–Lu hub edges are global. A node in block 3 can point at a hub that lives in block 500. To decode block 3 the decoder needs that hub-neighbor's *features* (its anchor embedding + structural features), but block 500 isn't in memory with the block-approach suggested above. + +**The fix: a parametric hub field.** The decoder doesn't care *which* node a hub is — it reads only the hub's **anchor** (for the embedding `cluster_mlp(anchor)`) and its **structural feature** `[log1p(in_deg), log1p(mean_nbr_in_deg)]`. The second term (a hub's neighbors' mean in-degree) is almost constant across hubs because every node's `k` neighbors are drawn the same way (~`knn_frac` coherent + ~`(1-knn_frac)` other hubs). So we hold it at one analytic value and let only the **anchor** and **in-degree** vary. That leaves just two properties determining a hub, so we store no hub nodes at all and sample those two per Chung–Lu edge: + - **`anchor ~ uniform`** over the `SS` anchors (weight and anchor are independent in the blend, so a target drawn with probability proporional to the weight doesn't bias the anchor). + - **`in-degree ~ size-biased(indeg_dist)`**: targets are drawn proportional to the weight, so a selected hub's weight follows `P(w) ∝ w·hist(w)`. This is precomputed once from the sample's in-degree histogram. + +The default pipeline still uses the whole-graph path; pass `--block-local` to switch to this streaming generator. + +### Result: Block Decode vs. the real reference (50K→10M) + +Block Decode reproduces the real reference at 10M. (Real is the [§3](#3-headline-result) headline number). + +Build time: real **1,040 s** · Block Decode **1,021 s** + +| Search width `L` | Real 10M — Recall@10 (QPS) | Block Decode — Recall@10 (QPS) | +|:---:|:---:|:---:| +| 10 | 82.88% (31,387) | 82.42% (30,430) | +| 20 | 90.69% (19,719) | 92.54% (20,670) | +| 30 | 93.63% (16,958) | 96.16% (16,170) | +| 40 | 95.31% (14,493) | 97.66% (13,545) | +| 50 | 96.35% (12,617) | 98.53% (11,771) | +| 100 | 98.22% (7,481) | 99.80% (7,341) | +| 200 | 99.24% (4,527) | 100.00% (4,341) | +| 300 | 99.50% (3,255) | 100.00% (3,086) | + +To reproduce the block decode result: +``` +# 50K -> 10M +python generate_data.py --sample SAMPLE_50K.fbin --out-dir out/50k_10m \ + --target 9990000 --nc 500 --resid-rank 32 \ + --resid-scale 1.8 --knn-frac 0.35 --model-cache out/model_cache \ + --block-local --block-size 1000000 +python build_vamana.py "$d/base.fbin" --num-vectors 9990000 --save # Build using cuVS vamana +``` + +Should be validated for larger data too. + +--- + +## 6. TODOs and Future Work + +- **Cheap ground truth.** Need to add support for the cheaper cluster-probe (IVF-style `nprobe`) GT. Need a per-cluster deterministic generator and validation on real. +- **Residual smoothness (LID).** Synthetic search is slightly *easier* than real at high `L`: synthetic vectors have lower intrinsic dimension than real (a property of the Gaussian residual). Raising `--resid-rank` fixes LID but collapses recall; a learned on-manifold residual is the open lead. +- **Broader validation.** Confirm the match holds across other search `k`, model sizes (`--d-emb/--hidden/--depth`), other indexes (CAGRA, HNSW, IVF-PQ), and build/search params (`R`, `graph_degree`). +- **Privacy.** Anchors *are* the real sample points, so the output can leak the sample. Workaround is to fit on a larger sample, cluster it, and use the centroids as anchors (an aggregate, not any single real point). diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/block_local.py b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/block_local.py new file mode 100644 index 0000000000..2393ca516e --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/block_local.py @@ -0,0 +1,279 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""block_local.py — streaming (block-local) version of generate + decode. + +A memory-efficient alternative to (generate_graph_knn + decode_graph), enabled by +generate_data.py's ``--block-local`` flag: nodes are laid out in anchor order and +generated/decoded one block at a time, so the (N,k) graph and O(N) feature arrays +never materialise. The hub tail uses the parametric hub field. See README §5 +("Block Decode mode") for the detailed explanation on the design. +""" + +from __future__ import annotations + +import numpy as np +import torch +from knn_decoder import STRUCT_FEAT_DIM, sample_norms_percentile +from tqdm import tqdm +from upsample import _est_coherence +from utils import normalize_features, sample_residuals, write_fbin_header + + +def _coherent_edges(loc, sz_of, gstart, k, w, rng): + """Windowed within-cluster wiring -> (b, k) in-block global neighbor ids.""" + half = max(1, w // 2) + offd = rng.integers(-half, half + 1, size=(loc.shape[0], k)) + tp = (loc[:, None] + offd) % sz_of[:, None] + tp = np.where(tp == loc[:, None], (tp + 1) % sz_of[:, None], tp) + return gstart[:, None] + tp + + +def _cluster_layout(nc, N, rng): + """Spread N nodes uniformly over the nc anchors, laid out in anchor order. + Returns (sizes (nc,), off (nc+1,)) int64 — off[c] = start id of anchor c. + """ + sizes = rng.multinomial(N, np.full(nc, 1.0 / nc)).astype(np.int64) + off = np.zeros(nc + 1, dtype=np.int64) + np.cumsum(sizes, out=off[1:]) + return sizes, off + + +def _hub_field_setup(indeg_dist, k, knn_frac): + """Precompute the on-the-fly hub sampler from the sample in-degree histogram. + Returns (sb_cdf, cl_scale, hub_mni): + + - `sb_cdf` : (SS,) float64 — inverse-CDF for the size-biased in-degree draw. + - `cl_scale`: scalar float for transforming weight -> expected Chung-Lu in-degree. Total + Chung-Lu edges `e_cl = N·k·(1-knn_frac)` are spread proportional to the weight + over all N nodes, so a node with weight w expects + `e_cl·w/(N·E[w]) = k·(1-knn_frac)·w/E[w]` edges. + `cl_scale = k·(1-knn_frac)/E[w]` (N cancels). Multiplied by the weight + to get an in-degree during the decode loop. + - `hub_mni` : scalar float — a hub's mean-neighbor-in-degree (approx): its k nbrs + are ~knn_frac coherent (low) and ~(1-knn_frac) other hubs (high, at + the size-biased mean in-degree). + """ + w = indeg_dist.astype(np.float64) + 1e-9 + mean_w = float(w.mean()) + cl_scale = float(k * (1.0 - knn_frac) / mean_w) + sb_cdf = np.cumsum(w) + sb_cdf /= sb_cdf[-1] + mean_sb_w = float( + (w * w).sum() / w.sum() + ) # E[w^2]/E[w] = size-biased mean weight + mean_sb_indeg = mean_sb_w * cl_scale + coh_deg = knn_frac * k + hub_mni = float(knn_frac * coh_deg + (1.0 - knn_frac) * mean_sb_indeg) + return sb_cdf, cl_scale, hub_mni + + +@torch.no_grad() +def generate_block_local( + model, + anchors, + mu, + sd, + feat_mu, + feat_sd, + resid_params, + norm_q, + cluster_of_anchor, + stats, + N, + n_queries, + k, + knn_frac, + seed, + device, + base_path, + block_size=1_000_000, +): + """Cluster-ordered, block-fused generate+decode. Streams base rows to `base_path` + and returns the held-out queries (n_queries, D). See module docstring. + + anchors : (nc, D) host anchor table (row = anchor = sample point). + """ + indeg_dist = stats["indeg_dist"] + nc = len(indeg_dist) # # anchors (= sample size) + tgt = stats.get("coherence", 0.3) or 0.3 + D = mu.shape[1] + w = max(k + 1, round(0.7 * k / max(tgt, 1e-3))) + + rng = np.random.default_rng(seed) + sizes, off = _cluster_layout(nc, N, rng) + + mu_t = torch.tensor(mu, device=device) + sd_t = torch.tensor(sd, device=device) + + sb_cdf, cl_scale, hub_mni_p = _hub_field_setup(indeg_dist, k, knn_frac) + + # held-out queries + q_idx = np.sort( + np.random.default_rng(seed + 999).choice( + N, size=n_queries, replace=False + ) + ) + base_count = N - n_queries + queries = np.empty((n_queries, D), dtype=np.float32) + q_off = 0 + fbase = open(base_path, "wb") + write_fbin_header(fbase, base_count, D) + + # group whole clusters into ~block_size-node blocks + blocks, c0, acc = [], 0, 0 + for c in range(nc): + acc += sizes[c] + if acc >= block_size or c == nc - 1: + blocks.append((c0, c + 1)) + c0, acc = c + 1, 0 + + # one proportional window correction (mirrors triadic) on the first block + ca0, cb0 = blocks[0] + lo0, hi0 = int(off[ca0]), int(off[cb0]) + if hi0 - lo0 > k + 1: + g0 = np.repeat(off[ca0:cb0], sizes[ca0:cb0]) + s0 = np.repeat(sizes[ca0:cb0], sizes[ca0:cb0]) + l0 = np.arange(lo0, hi0) - g0 + c0e = ( + _coherent_edges(l0, s0, g0, k, w, np.random.default_rng(seed + 7)) + - lo0 + ) + C = _est_coherence(c0e.astype(np.int64), seed=1) + if C > 0 and abs(C - tgt) > 0.015: + w = max(k + 1, round(w * C / tgt)) + + base_k = int(np.floor(knn_frac * k)) + fracp = knn_frac * k - base_k + ar_k = np.arange(k) + torch.manual_seed(seed + 1) + + for bi, (ca, cb) in enumerate( + tqdm(blocks, desc="block-local gen+decode", unit="block") + ): + lo, hi = int(off[ca]), int(off[cb]) + b = hi - lo + if b == 0: + continue + brng = np.random.default_rng(seed + 100 + bi) + + # --- per-node cluster / local-position --- + cl_of = np.repeat( + np.arange(ca, cb, dtype=np.int64), sizes[ca:cb] + ) # (b,) anchor id + gstart = np.repeat(off[ca:cb], sizes[ca:cb]) # (b,) anchor start id + sz_of = np.repeat(sizes[ca:cb], sizes[ca:cb]) # (b,) anchor size + loc = np.arange(lo, hi, dtype=np.int64) - gstart # (b,) local index + + # --- coherent (windowed within-cluster) neighbors --- + coh_gid = _coherent_edges( + loc, sz_of, gstart, k, w, brng + ) # (b, k) in [lo, hi) + coh_loc = coh_gid - lo # (b, k) local + + # --- blend mask: k_coh coherent, the rest Chung-Lu --- + k_coh = np.clip( + base_k + (brng.random(b) < fracp).astype(np.int64), 0, k + ) + mask = ar_k[None, :] < k_coh[:, None] # (b, k) True=coherent + + # --- Chung-Lu targets: sample (anchor, in-degree) per edge --- + t_anchor = brng.integers(0, nc, size=(b, k)).astype( + np.int64 + ) # anchor ~ uniform + t_w = indeg_dist[ + np.searchsorted(sb_cdf, brng.random((b, k))) + ] # in-degree ~ size-biased + t_indeg = t_w.astype(np.float32) * cl_scale # (b, k) + + # --- in-degree: coherent counted locally + own expected Chung-Lu in-degree --- + node_indeg = np.bincount(coh_loc[mask], minlength=b).astype(np.float32) + node_w = indeg_dist[brng.integers(0, nc, size=b)].astype( + np.float32 + ) # own weight ~ uniform + node_indeg += node_w * cl_scale + + # neighbor in-degree: coherent -> node_indeg[target]; Chung-Lu -> sampled in-degree + nbr_indeg = np.where(mask, node_indeg[coh_loc], t_indeg) # (b, k) + mni = nbr_indeg.mean(1) # (b,) + + # --- structural features (node + neighbors), normalized as in training --- + struct_node = np.stack( + [np.log1p(node_indeg), np.log1p(mni)], 1 + ) # (b, 2) + hub_nbr_struct = np.stack( + [ + np.log1p(t_indeg), + np.full((b, k), np.log1p(hub_mni_p), np.float32), + ], + -1, + ) + struct_nbr = np.where( + mask[:, :, None], struct_node[coh_loc], hub_nbr_struct + ) + struct_node = normalize_features(struct_node, feat_mu, feat_sd) + struct_nbr = normalize_features( + struct_nbr.reshape(b * k, STRUCT_FEAT_DIM), feat_mu, feat_sd + ).reshape(b, k, STRUCT_FEAT_DIM) + + cluster_of_node = cluster_of_anchor[cl_of].astype(np.int64) # (b,) + + # --- decode: stream this block's anchors, MLP them --- + block_emb = model.cluster_mlp( + torch.as_tensor( + np.ascontiguousarray(anchors[ca:cb]), device=device + ) + ) # (b_clusters, d_emb) + node_emb = block_emb[ + torch.tensor(cl_of - ca, device=device) + ] # (b, d_emb) + + # coherent nbr shares the node's anchor emb; Chung-Lu nbr -> MLP of its + # sampled anchor (gathered from the host anchor table, this block only). + nbr_emb = ( + node_emb[:, None, :].expand(-1, k, -1).contiguous() + ) # (b, k, d_emb) + is_cl = ~mask + if is_cl.any(): + cl_anchor = t_anchor[is_cl] # (n_cl,) + cl_rows = torch.as_tensor( + np.ascontiguousarray(anchors[cl_anchor]), device=device + ) # (n_cl, D) + nbr_emb[torch.as_tensor(is_cl, device=device)] = model.cluster_mlp( + cl_rows + ) + + hi_f = torch.cat( + [node_emb, torch.as_tensor(struct_node, device=device)], -1 + ) + hj_f = torch.cat( + [nbr_emb, torch.as_tensor(struct_nbr, device=device)], -1 + ) + pred = model(hi_f, hj_f) + pred = pred + sample_residuals( + torch.tensor(cluster_of_node, device=device), + resid_params, + D, + device, + ) + xb = pred * sd_t + mu_t + norm = sample_norms_percentile(cluster_of_node, norm_q, seed=seed + bi) + dirv = xb / (xb.norm(dim=1, keepdim=True) + 1e-12) + xb = dirv * torch.tensor(norm, device=device).unsqueeze(1) + xb_np = xb.detach().cpu().numpy() + + # --- split base / queries: this block's query ids --- + qa, qb = np.searchsorted(q_idx, lo), np.searchsorted(q_idx, hi) + qloc = q_idx[qa:qb] - lo + if qloc.size: + qm = np.zeros(b, dtype=bool) + qm[qloc] = True + np.ascontiguousarray(xb_np[~qm]).tofile(fbase) + queries[q_off : q_off + qloc.size] = xb_np[qloc] + q_off += qloc.size + else: + xb_np.tofile(fbase) + + fbase.close() + return queries diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/build_vamana.py b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/build_vamana.py new file mode 100644 index 0000000000..bf9ee864a6 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/build_vamana.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build a cuVS Vamana index on a subset of an fbin dataset. + +Reads an .fbin file, takes the first N vectors, moves them to the GPU +and runs cuVS Vamana with the given parameters while timing the build. + +""" + +import argparse +import os +import time +from contextlib import contextmanager + +import numpy as np + + +@contextmanager +def timer(label): + """Context manager that prints the wall-clock time spent in the block.""" + start = time.perf_counter() + try: + yield + finally: + elapsed = time.perf_counter() - start + print(f"[timer] {label}: {elapsed:.3f} s") + + +def read_fbin_subset(path, num_vectors, header_dtype=np.int32): + """Read the first ``num_vectors`` vectors from an .fbin file. + + The .fbin layout is a 2-element header ``[n_vectors, dim]`` followed by + ``n_vectors * dim`` float32 values. Standard .fbin uses int32 headers; + pass ``np.int64`` for .fbin64-style files. + """ + header_bytes = header_dtype().itemsize * 2 + with open(path, "rb") as f: + header = np.fromfile(f, count=2, dtype=header_dtype) + total_vectors, dim = int(header[0]), int(header[1]) + + vectors_to_read = min(num_vectors, total_vectors) + print( + f"File: {total_vectors:,} vectors, dim={dim}. " + f"Reading first {vectors_to_read:,} vectors." + ) + + # Header is already consumed; read only the needed float32 block. + f.seek(header_bytes) + data = np.fromfile(f, count=vectors_to_read * dim, dtype=np.float32) + + actual = data.size // dim + if actual < vectors_to_read: + print( + f"Warning: only read {actual:,} vectors instead of {vectors_to_read:,}" + ) + vectors_to_read = actual + + data = data[: vectors_to_read * dim].reshape(vectors_to_read, dim) + return data, dim + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("fbin", help="Path to the input .fbin file") + parser.add_argument( + "--num-vectors", + type=int, + default=10_000_000, + help="Number of vectors from the start of the file to use (default: 10M)", + ) + parser.add_argument("--max-fraction", type=float, default=0.06) + parser.add_argument("--visited-size", type=int, default=256) + parser.add_argument("--graph-degree", type=int, default=64) + parser.add_argument("--vamana-iters", type=int, default=1) + parser.add_argument("--alpha", type=float, default=1.2) + parser.add_argument( + "--metric", + default="sqeuclidean", + help="Distance metric (default: sqeuclidean)", + ) + parser.add_argument( + "--header-dtype", + choices=["int32", "int64"], + default="int32", + help="dtype of the 2-element fbin header (default: int32)", + ) + parser.add_argument( + "--save", + action="store_true", + help="Save the built index as vamana.index in the input .fbin's directory", + ) + args = parser.parse_args() + + # Imported here so the script can at least show --help without cuVS/cupy. + import cupy as cp + from cuvs.neighbors import vamana + + header_dtype = np.int32 if args.header_dtype == "int32" else np.int64 + + with timer("read fbin subset (host)"): + host_data, dim = read_fbin_subset( + args.fbin, args.num_vectors, header_dtype + ) + + n = host_data.shape[0] + gb = host_data.nbytes / 1e9 + print(f"Loaded {n:,} x {dim} float32 ({gb:.2f} GB) on host.") + + with timer("copy dataset host -> device"): + device_data = cp.asarray(host_data) + cp.cuda.Stream.null.synchronize() + + index_params = vamana.IndexParams( + metric=args.metric, + graph_degree=args.graph_degree, + visited_size=args.visited_size, + vamana_iters=args.vamana_iters, + alpha=args.alpha, + max_fraction=args.max_fraction, + ) + + print( + "Vamana params: " + f"metric={args.metric}, graph_degree={args.graph_degree}, " + f"visited_size={args.visited_size}, vamana_iters={args.vamana_iters}, " + f"alpha={args.alpha}, max_fraction={args.max_fraction}" + ) + + with timer(f"cuVS Vamana build ({n:,} vectors)"): + index = vamana.build(index_params, device_data) + cp.cuda.Stream.null.synchronize() + + print("Build complete.") + + if args.save: + save_path = os.path.join( + os.path.dirname(os.path.abspath(args.fbin)), "vamana.index" + ) + with timer("save index"): + vamana.save(save_path, index) + print(f"Saved index to {save_path}") + + +if __name__ == "__main__": + main() diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/figures/results_comparison.png b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/figures/results_comparison.png new file mode 100644 index 0000000000..ecaf137da1 Binary files /dev/null and b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/figures/results_comparison.png differ diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/figures/workflow.png b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/figures/workflow.png new file mode 100644 index 0000000000..e306353750 Binary files /dev/null and b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/figures/workflow.png differ diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/generate_data.py b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/generate_data.py new file mode 100644 index 0000000000..a5cc93b8a1 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/generate_data.py @@ -0,0 +1,474 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""generate_data.py - entry point: fit on a real sample, generate a synthetic benchmark bundle.""" + +from __future__ import annotations + +import argparse +import gc +import os +import time + +import numpy as np +import torch +from block_local import generate_block_local +from knn_decoder import ( + decode_graph, + fit_norm_quantiles, + fit_residuals, + kNNDecoder, + sample_norms_percentile, + train_model, +) +from upsample import _est_coherence, generate_graph_knn +from utils import ( + PhaseTimer, + build_all_neighbors, + build_kmeans, + holdout_split, + load_fbin, + section, + step, + write_bundle, + write_bundle_streamed, +) + + +def get_args(): + p = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + # ================================================================== # + # DATA / RUN — input sample, output, and what to generate + # ================================================================== # + g_data = p.add_argument_group( + "data / run", "input sample, output bundle, and what to generate" + ) + g_data.add_argument( + "--sample", + type=str, + required=True, + help="path to real .fbin sample to fit on.", + ) + g_data.add_argument( + "--out-dir", + type=str, + required=True, + help="dir to write base.fbin / queries.fbin / groundtruth.*", + ) + g_data.add_argument( + "--target", + type=int, + default=None, + help="number of synthetic BASE vectors to generate. If this is not given, " + "we just split the sample data to base and n_queries and return.", + ) + g_data.add_argument( + "--n-queries", + type=int, + default=10000, + help="held-out queries, always disjoint from base", + ) + g_data.add_argument( + "--gt-k", + type=int, + default=100, + help="exact GT neighbors for each query", + ) + g_data.add_argument( + "--seed", + type=int, + default=42, + help="global RNG seed (KMeans, graph gen, decode)", + ) + + # ================================================================== # + # kNN UPSAMPLING: build the N-node kNN graph from the sample's kNN + # ================================================================== # + g_up = p.add_argument_group( + "kNN upsampling (small real kNN -> large synthetic kNN)" + ) + g_up.add_argument( + "--knn", + type=int, + default=10, + help="k for the kNN of the sample training graph AND the generated " + "graph for decoding. Don't change this unless you know what you are doing.", + ) + g_up.add_argument( + "--knn-frac", + type=float, + default=0.66, + help="fraction of each node's edges from the synthesized kNN; the " + "rest are random. 1.0=fully coherent, lower=harder.", + ) + + # ================================================================== # + # kNN DECODING: the model that maps the kNN graph -> vectors + # ================================================================== # + g_dec = p.add_argument_group( + "kNN decoding (kNN graph -> vector embeddings)" + ) + g_dec.add_argument( + "--d-emb", type=int, default=64, help="cluster-ID embedding dimension" + ) + g_dec.add_argument( + "--hidden", type=int, default=1024, help="Model MLP hidden width" + ) + g_dec.add_argument("--depth", type=int, default=3, help="Model MLP layers") + g_dec.add_argument( + "--epochs", type=int, default=400, help="Model training epochs" + ) + g_dec.add_argument( + "--lr", type=float, default=5e-3, help="Model learning rate" + ) + g_dec.add_argument( + "--batch", + type=int, + default=1024, + help="Model training batch size for SGD", + ) + g_dec.add_argument( + "--nc", + type=int, + default=None, + help="# clusters for the per-cluster residual/norm fit — same as " + "cuvs-bench synthesize_dataset's `nc`. Recommended default: sample_size / 100.", + ) + g_dec.add_argument( + "--resid-rank", + type=int, + default=64, + help="rank of the per-cluster residual Gaussian " + "(within-cluster spread added to the Model output). This is similar to" + "pca-components in cuvs_bench.", + ) + g_dec.add_argument( + "--resid-scale", + type=float, + default=1.0, + help="scale the residual spread (0=Model mean only)", + ) + g_dec.add_argument( + "--model-cache", + type=str, + default=None, + help="dir to cache the full fit bundle (model + mu/sd/feat-stats + " + "resid_params + norm_q + cluster_ids + stats). A cache hit skips " + "the whole fit and reproduces the dataset: identical bundle + " + "identical generate args (knn-frac/target/seed) => identical data.", + ) + g_dec.add_argument( + "--host-gather", + action="store_true", + help="hold the training anchor table / kNN graph / struct on the " + "HOST and gather each minibatch to the GPU per step. Use only when the sample " + "is too big to fit in GPU memory (100M+ sample).", + ) + + # ================================================================== # + # BLOCK-LOCAL: cluster-ordered streaming generate+decode + # ================================================================== # + g_bl = p.add_argument_group("block-local (scaling path — see README §5)") + g_bl.add_argument( + "--block-local", + action="store_true", + help="use the cluster-ordered block-local generator (block_local.py) instead " + "of the default whole-graph pool+decode. Streams one cluster-ordered block " + "at a time. This option allows to never materialize the full (N,k) graph or the " + "O(N) feature arrays. USE WHEN: the target (N, k) graph is too large for the " + "default path to hold in memory (roughly multi-billion-scale+); the " + "whole-graph path is simpler and fine for smaller ones (1B and under).", + ) + g_bl.add_argument( + "--block-size", + type=int, + default=1_000_000, + help="[--block-local] target # nodes per streamed block — caps peak " + "GPU memory. Default 1M is okay, but lower it if you hit GPU OOM, raise it to " + "cut per-block overhead when memory allows.", + ) + + return p.parse_args() + + +def main(): + args = get_args() + device = "cuda" if torch.cuda.is_available() else "cpu" + os.makedirs(args.out_dir, exist_ok=True) + timer = PhaseTimer() # accumulates fit/generate phase times + + # ------------------------------------------------------------------ # + # CONFIG + LOAD + # ------------------------------------------------------------------ # + is_synth = args.target is not None + + section("CONFIG") + print(f" sample: {args.sample}") + print(f" out_dir: {args.out_dir}") + print(f" run: {'SYNTH' if is_synth else 'REAL (no target)'}") + print(f" target: {args.target} n_queries: {args.n_queries}") + + section("LOAD DATA") + t0 = time.perf_counter() + X = load_fbin(args.sample) + D = X.shape[1] + print(f" shape: {X.shape}", flush=True) + timer.lap("fit", "load data", t0) + + # ------------------------------------------------------------------ # + # REAL reference (no --target): just split the sample and stop + # ------------------------------------------------------------------ # + if not is_synth: + section("REAL — split sample into base + held-out queries") + base, queries = holdout_split(X, args.n_queries, args.seed) + write_bundle(args.out_dir, base, queries, args.gt_k) + section("DONE (just split the sample)") + return + + # ------------------------------------------------------------------ # + # SYNTH: FIT — model + residual + norm + clusters + stats. + # ------------------------------------------------------------------ # + fit_path = None + if args.model_cache: + os.makedirs(args.model_cache, exist_ok=True) + samp = os.path.splitext(os.path.basename(args.sample))[0][:32] + req_nc = args.nc if args.nc else max(1, len(X) // 100) + key = ( + f"fitb_{samp}_ss{len(X)}_emb{args.d_emb}_h{args.hidden}_d{args.depth}" + f"_knn{args.knn}_ep{args.epochs}_nc{req_nc}_rr{args.resid_rank}" + f"_rs{args.resid_scale}_s{args.seed}" + ) + fit_path = os.path.join(args.model_cache, key + ".pt") + + if fit_path and os.path.exists(fit_path): + section("LOAD FIT BUNDLE (cached — skipping fit)") + t0 = time.perf_counter() + step(f"loading cached fit bundle: {fit_path}") + ck = torch.load(fit_path, map_location=device, weights_only=False) + model = kNNDecoder( + d_out=D, d_emb=args.d_emb, hidden=args.hidden, depth=args.depth + ).to(device) + model.load_state_dict(ck["model"]) + model.eval() + mu, sd, feat_mu, feat_sd = ( + ck["mu"], + ck["sd"], + ck["feat_mu"], + ck["feat_sd"], + ) + resid_params = ck["resid_params"] + norm_q = ck["norm_q"] + cluster_ids = ck["cluster_ids"] + stats = ck["stats"] + nc = ck["nc"] + print( + f" model + resid_params + norm_q + cluster_ids + stats loaded " + f"(nc={nc}) — fit reproduced from cache", + flush=True, + ) + timer.lap("fit", "load fit bundle", t0) + else: + section("kNN Graph on sample") + t0 = time.perf_counter() + step(f"all-neighbors kNN (k={args.knn}) on {len(X):,} pts ...") + sample_knn = build_all_neighbors(X, args.knn) + timer.lap("fit", "sample kNN", t0) + + section("KMeans on sample") + t0 = time.perf_counter() + ss = len(X) + nc = args.nc if args.nc else max(1, ss // 100) + cluster_ids, _ = build_kmeans(X, nc, seed=args.seed) + nc = int(cluster_ids.max()) + 1 + timer.lap("fit", "kmeans", t0) + + section("Measuring kNN graph statistics") + t0 = time.perf_counter() + stats = { + "coherence": _est_coherence( + sample_knn, seed=args.seed + ), # float value + "indeg_dist": np.bincount( + sample_knn.reshape(-1), minlength=ss + ).astype(np.float32), # (ss, ) + } + print( + f" coherence={stats['coherence']:.4f} " + f"in-deg max={stats['indeg_dist'].max():.0f}", + flush=True, + ) + timer.lap("fit", "measure stats", t0) + + section("TRAIN MODEL") + t0 = time.perf_counter() + model, mu, sd, feat_mu, feat_sd, struct_feats = train_model( + X, + sample_knn, + stats["indeg_dist"], + d_emb=args.d_emb, + hidden=args.hidden, + depth=args.depth, + epochs=args.epochs, + lr=args.lr, + batch=args.batch, + device=device, + seed=args.seed, + host_gather=args.host_gather, + ) + model.eval() + timer.lap("fit", "train model", t0) + + section("FIT RESIDUAL MODEL") + t0 = time.perf_counter() + resid_params = fit_residuals( + X, + sample_knn, + model, + mu, + sd, + feat_mu, + feat_sd, + nc=nc, + rank=args.resid_rank, + device=device, + scale=args.resid_scale, + struct_feats=struct_feats, + resid_ids=cluster_ids, + in_deg=stats["indeg_dist"], + ) + timer.lap("fit", "fit residual", t0) + + section("FIT NORM (radial percentile scheme)") + t0 = time.perf_counter() + norm_q, nmean, ncv = fit_norm_quantiles(X, cluster_ids, nc) + print( + f" real norms: mean={nmean:.4f} cv={ncv:.4f} " + f"-> percentile (per-cluster norm inverse-CDF)", + flush=True, + ) + timer.lap("fit", "fit norm", t0) + + if fit_path: + torch.save( + { + "model": model.state_dict(), + "mu": mu, + "sd": sd, + "feat_mu": feat_mu, + "feat_sd": feat_sd, + "resid_params": resid_params, + "norm_q": norm_q, + "cluster_ids": cluster_ids, + "stats": stats, + "nc": nc, + }, + fit_path, + ) + print(f" cached fit bundle -> {fit_path}", flush=True) + + # ------------------------------------------------------------------ # + # Build the synthetic kNN graph, decode, split + # ------------------------------------------------------------------ # + n_pool = args.target + args.n_queries # base (=target) + held-out queries + base_path = os.path.join(args.out_dir, "base.fbin") + + if args.block_local: + # cluster-ordered streaming generate+decode. Fuses graph gen and decode + # per block, so the (N,k) graph and O(N) features never materialize. + section("BLOCK-LOCAL generate+decode (parametric hub field)") + t0 = time.perf_counter() + queries = generate_block_local( + model, + X, + mu, + sd, + feat_mu, + feat_sd, + resid_params, + norm_q, + cluster_ids, + stats, + N=n_pool, + n_queries=args.n_queries, + k=args.knn, + knn_frac=args.knn_frac, + seed=args.seed, + device=device, + base_path=base_path, + block_size=args.block_size, + ) + print( + f" generated {n_pool:,} pts (base streamed to {base_path})", + flush=True, + ) + timer.lap("generate", "block-local gen+decode", t0) + else: + section("BUILD POOL GRAPH") + t0 = time.perf_counter() + anchor_pool, nbr_pool = generate_graph_knn( + stats, N=n_pool, k=args.knn, seed=args.seed, knn_frac=args.knn_frac + ) + print(f" pool graph: {len(anchor_pool):,} nodes", flush=True) + timer.lap("generate", "pool graph", t0) + + section("DECODE kNN graph — stream base to disk") + cluster_pool = cluster_ids[ + anchor_pool + ] # anchor -> cluster for residual + norm + norm_target = sample_norms_percentile( + cluster_pool, norm_q, seed=args.seed + ) + q_idx = np.sort( + np.random.default_rng(args.seed).choice( # held-out query ids + n_pool, min(args.n_queries, n_pool), replace=False + ) + ) + t0 = time.perf_counter() + # Streaming decode: base rows go straight to base.fbin (host RAM ~ one chunk), + # only the held-out queries come back in memory. + queries = decode_graph( + anchor_pool, + nbr_pool, + model, + X, + mu, + sd, + feat_mu, + feat_sd, + resid_params, + norm_target, + device, + args.seed, + resid_ids=cluster_pool, + base_path=base_path, + query_idx=q_idx, + ) + print( + f" decoded {n_pool:,} pts (base streamed to {base_path})", + flush=True, + ) + timer.lap("generate", "decode", t0) + + model.to("cpu") + del resid_params + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + + section("WRITE queries + exact GT") + t0 = time.perf_counter() + write_bundle_streamed( + args.out_dir, base_path, n_pool - len(queries), queries, args.gt_k, D + ) + timer.lap("generate", "write + GT", t0) + + timer.summary() + section("DONE (synth)") + + +if __name__ == "__main__": + main() diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/knn_decoder.py b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/knn_decoder.py new file mode 100644 index 0000000000..8892d106e7 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/knn_decoder.py @@ -0,0 +1,572 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""kNN decoder: map a graph -> vectors (mean + per-cluster residual).""" + +from __future__ import annotations + +import threading +import time + +import numpy as np +import torch +from torch import nn +from tqdm import tqdm +from utils import ( + CHUNK_SIZE, + batched_mean_std, + normalize_features, + sample_residuals, + ts, + write_fbin_header, +) + +STRUCT_FEAT_DIM = 2 +_NORM_QUANTILE_COUNT = ( + 256 # per-cluster norm inverse-CDF grid (synthesize_dataset/_fit.py) +) + + +def mlp(d_in, d_out, hidden, depth=2): + """A simple GELU MLP: d_in -> [hidden]*depth -> d_out.""" + layers = [nn.Linear(d_in, hidden), nn.GELU()] + for _ in range(depth - 1): + layers += [nn.Linear(hidden, hidden), nn.GELU()] + layers.append(nn.Linear(hidden, d_out)) + return nn.Sequential(*layers) + + +def compute_structural_features(nbr, in_deg=None): + """Per-node topological features from a kNN graph: [log1p(in-degree), + log1p(mean neighbour in-degree)]. Returns (n_nodes, STRUCT_FEAT_DIM) float32. + + in_deg : optional (n_nodes,) per-node in-degree to reuse. + """ + n_nodes = nbr.shape[0] + if in_deg is None: + in_deg = np.bincount(nbr.reshape(-1), minlength=n_nodes) + mean_nbr_indeg = in_deg[nbr].mean(axis=1).astype(np.float32) # (n_nodes,) + return np.stack( + [np.log1p(in_deg), np.log1p(mean_nbr_indeg)], axis=1 + ).astype(np.float32) # (n_nodes, F) + + +class kNNDecoder(nn.Module): + """Map local graph structure (anchor + connectivity) -> ambient vector.""" + + def __init__( + self, + d_out=1024, + d_emb=64, + d_struct=STRUCT_FEAT_DIM, + hidden=1024, + depth=3, + ): + super().__init__() + self.d_struct = d_struct + + # A node's embedding is a function of its anchor (the raw sample point). + self.cluster_mlp = mlp(d_out, d_emb, 256, depth=2) # anchor -> d_emb + + d_node = d_emb + d_struct # total node feature dimension + + # Score each neighbor relative to the centre node + self.att = nn.Linear(d_node * 2, 1) + + # Decode [h_i, aggregate_neighbors] -> embedding + self.out = mlp(d_node * 2, d_out, hidden, depth) + + def node_feat(self, anchor_vecs, struct): + """Build node feature vectors from anchor rows. + + anchor_vecs : (B, d_out) float32 — the raw anchor rows + struct : (B, d_struct) float32 — normalized topological features + """ + c = self.cluster_mlp(anchor_vecs) # (B, d_emb) + return torch.cat([c, struct], dim=-1) # (B, d_node) + + def forward(self, h_i, h_nbrs): + """Aggregate neighbors and project to ambient space. + + h_i : (B, d_node) — centre node features + h_nbrs : (B, k, d_node) — neighbor features + -> (B, d_out) + """ + B, k, d = h_nbrs.shape + hi_exp = h_i.unsqueeze(1).expand(B, k, d) # (B, k, d_node) + scores = self.att(torch.cat([hi_exp, h_nbrs], -1)) # (B, k, 1) + w = torch.softmax(scores, dim=1) # (B, k, 1) + agg = (w * h_nbrs).sum(1) # (B, d_node) + return self.out(torch.cat([h_i, agg], -1)) # (B, d_out) + + +def train_model( + X, + sample_knn, + in_deg, + d_emb, + hidden, + depth, + epochs, + lr, + batch, + device, + seed, + host_gather=False, +): + """Train kNNDecoder to map graph structure -> embedding. + + The sample points X are used as the anchors (row i = anchor i) and as the + regression targets. + + Parameters + ---------- + X : (ss, D) float32 real sample embeddings (ss = sample size); used as + the anchors and the training targets. + sample_knn : (ss, k) int64 all-neighbors kNN graph of the sample + in_deg : (ss,) per-node in-degree of the sample kNN (the stats' indeg_dist), + reused for the structural features. + host_gather: if False (default), the anchor table / kNN graph / struct features are + held GPU-RESIDENT and gathered on-device. This is fast, for samples that + fit in GPU memory. If True, they stay on the HOST and each minibatch's + rows are shipped to the GPU per step — slower per step but scales to + samples too big for GPU memory. Same loop either way: only the storage + device changes (the per-batch `.to(device)` is a no-op when resident). + + + Returns (model, mu, sd, feat_mu, feat_sd, struct_feats): + mu/sd per-dim stats of X + feat_mu/feat_sd stats of the structural features, reused at generate time. + """ + torch.manual_seed(seed) + + ss, D = X.shape + k = sample_knn.shape[1] + + mu, sd = batched_mean_std(X) + + # Structural node features from the real kNN graph + struct_feats = compute_structural_features(sample_knn, in_deg) + feat_mu, feat_sd = batched_mean_std(struct_feats) + struct_feats = normalize_features(struct_feats, feat_mu, feat_sd) + + mu_t = torch.as_tensor(mu, device=device) # (1, D) for per-batch normalize + sd_t = torch.as_tensor(sd, device=device) + + store = "cpu" if host_gather else device + Xs = torch.as_tensor(X, device=store) # (ss, D) anchors == targets + nbrs = torch.as_tensor(sample_knn, device=store).long() # (ss, k) + structs = torch.as_tensor(struct_feats, device=store) # (ss, F) + + model = kNNDecoder(d_out=D, d_emb=d_emb, hidden=hidden, depth=depth).to( + device + ) + opt = torch.optim.Adam(model.parameters(), lr=lr) + + n_batches = (ss + batch - 1) // batch + print( + f" [{ts()}] training on {ss:,} nodes, {n_batches:,} batches/epoch, {epochs} " + f"epochs ({'host-gather' if host_gather else 'GPU-resident'}) ...", + flush=True, + ) + t0 = time.perf_counter() + + for ep in range(epochs): + perm = torch.randperm(ss, device=store) # on `store` + total_loss = 0.0 + + for bi, i in enumerate(range(0, ss, batch)): + idx = perm[i : i + batch] # on `store` + b = idx.numel() + + xi = ( + Xs[idx].to(device) - mu_t + ) / sd_t # (B, D); .to() no-op if resident + + flat = nbrs[idx].reshape(-1) # (B*k,) + allids = torch.cat([idx, flat]) # (B + B*k,) + uniq, inv = torch.unique(allids, return_inverse=True) + a_uniq = Xs[uniq].to(device) # (U, D) unique anchor rows + s_uniq = structs[uniq].to(device) + ufeat = model.node_feat(a_uniq, s_uniq) # cluster_mlp once/unique + inv = inv.to(device) # gather the device-side ufeat + + hi = ufeat[inv[:b]] # (B, d_node) + hj = ufeat[inv[b:]].reshape(b, k, -1) # (B, k, d_node) + + x_pred = model(hi, hj) # (B, D) + loss = ((x_pred - xi) ** 2).mean() + + opt.zero_grad() + loss.backward() + opt.step() + + total_loss += loss.item() * b + + if ( + ep == 0 + and n_batches > 200 + and bi > 0 + and bi % (n_batches // 5) == 0 + ): + print( + f" [{ts()}] epoch 0: batch {bi:,}/{n_batches:,} " + f"({100 * bi / n_batches:.0f}%)", + flush=True, + ) + + ep_loss = total_loss / ss + + if ep % 10 == 0 or ep == epochs - 1: + elapsed = time.perf_counter() - t0 + print( + f" [{ts()}] ep {ep:4d}/{epochs} loss={ep_loss:.5f} " + f"elapsed={elapsed:.0f}s", + flush=True, + ) + + print(f" training done in {time.perf_counter() - t0:.0f}s", flush=True) + + return model, mu, sd, feat_mu, feat_sd, struct_feats + + +@torch.no_grad() +def fit_residuals( + X, + sample_knn, + model, + mu, + sd, + feat_mu, + feat_sd, + nc, + rank, + device, + scale=1.0, + chunk=CHUNK_SIZE, + struct_feats=None, + resid_ids=None, + in_deg=None, +): + """Fit per-cluster low-rank Gaussian on residuals x - GNN(x). + + Residuals are computed in the model's Z-normalized output space + ((x - mu) / sd). Returns a list of nc dicts (or None for tiny clusters), + each with GPU tensors {mean (D,), comps (r, D), stds (r,), noise_std}. + + Parameters + ---------- + X : (ss, D) float32 real sample embeddings (ss = sample size) — + the targets whose residual (x - decode(x)) we model. + sample_knn : (ss, k) int64 kNN graph, fed to the decoder to produce the + per-node mean GNN(x) that the residual is taken against. + model : the trained kNNDecoder — run (no-grad) to get the vector embeddings. + mu, sd : (1, D) float32 per-dim stats of X; the residual is fit in the + z-normalized space (x - mu) / sd (the space the model predicts). + feat_mu, : normalization stats for the structural features, reused here so + feat_sd the decode pass sees the exact same feature scale as training. + nc : int number of residual groups to fit == number of distinct + resid_ids (= # clusters). + rank : int target rank r of the per-cluster low-rank Gaussian + (clamped down to the cluster's point count). + device : torch device. + scale : shrink/scale the sampled residual spread. 1.0 = full real + within-cluster spread; 0.0 = GNN mean only (collapsed); + >1.0 roughens (the --resid-scale difficulty knob) the manifold. + chunk : batch size for the decode pass (caps peak memory). + struct_feats : precomputed structural features from train_model, + passed in to avoid recomputing them. + resid_ids : (ss,) int32 the KMeans cluster label per node. + """ + ss, D = X.shape + k = sample_knn.shape[1] + + if ( + struct_feats is None + ): # reused from train_model unless we loaded a cached model + struct_feats = compute_structural_features(sample_knn, in_deg) + struct_feats = normalize_features(struct_feats, feat_mu, feat_sd) + + # Decode all real nodes -> prediction; residual = target - prediction. + # For each chunk, gather the unique anchor rows (centers + neighbors) host->GPU + # and run cluster_mlp once each (Row i is anchor i, so the chunk's center ids are i..end.) + R = np.empty((ss, D), dtype=np.float32) + for i in tqdm( + range(0, ss, chunk), desc="residual: decode sample", unit="chunk" + ): + end = min(i + chunk, ss) + b = end - i + idx = np.arange(i, end) # center ids == rows (host) + flat = sample_knn[idx].reshape(-1) # (b*k,) neighbor ids (host) + allids = np.concatenate([idx, flat]) + uniq, inv = np.unique(allids, return_inverse=True) + a_uniq = torch.as_tensor( + X[uniq], device=device + ) # (U, D) anchor rows host->GPU + s_uniq = torch.as_tensor(struct_feats[uniq], device=device) + ufeat = model.node_feat(a_uniq, s_uniq) # MLP once/unique + inv = torch.as_tensor(inv, device=device) + + hi = ufeat[inv[:b]] # (b, d_node) + hj = ufeat[inv[b:]].reshape(b, k, -1) # (b, k, d_node) + pred = model(hi, hj) # normalized space + xt_chunk = ((X[i:end] - mu) / sd).astype( + np.float32 + ) # normalize this chunk only + R[i:end] = xt_chunk - pred.cpu().numpy() + + # Per-cluster PPCA fit + params = [None] * nc + for c in tqdm( + range(nc), desc="residual: per-cluster PPCA", unit="cluster" + ): + sel = np.where(resid_ids == c)[0] + sz = len(sel) + if sz < 2: + continue + Rc = torch.tensor(R[sel], device=device) # (sz, D) + mean_c = Rc.mean(0) # (D,) + Rc0 = Rc - mean_c + + # SVD: Vh rows are principal directions + _, S, Vh = torch.linalg.svd(Rc0, full_matrices=False) + r = int(min(rank, Vh.shape[0])) + comps = Vh[:r].contiguous() # (r, D) + var = (S**2) / max(sz - 1, 1) # per-direction variance + stds = torch.sqrt(var[:r]) * scale # (r,) — anisotropic sheet + noise_std = torch.zeros((), device=device) + params[c] = { + "mean": mean_c * scale, + "comps": comps, + "stds": stds, + "noise_std": noise_std, + } + + n_fit = sum(1 for p in params if p is not None) + print( + f" fit residual Gaussians for {n_fit}/{nc} clusters " + f"(rank<={rank}, scale={scale});", + flush=True, + ) + return params + + +def fit_norm_quantiles(X, cluster_ids, nc): + """Per-cluster norm inverse-CDF grids. + + Mirrors synthesize_dataset/_fit.py. Returns + (norm_quantiles: (nc, 256) float32, mean: float, cv: float). + Empty clusters fall back to the global grid. + """ + import cupy as cp + + n = len(X) + norms = cp.empty(n, dtype=cp.float32) # (ss,) on device + for s in range(0, n, CHUNK_SIZE): + e = min(s + CHUNK_SIZE, n) + chunk_gpu = cp.asarray(X[s:e], dtype=cp.float32) + norms[s:e] = cp.linalg.norm(chunk_gpu, axis=1) + del chunk_gpu + mean = float(norms.mean()) + cv = float(norms.std() / max(mean, 1e-12)) + levels = cp.linspace(0.0, 1.0, _NORM_QUANTILE_COUNT) + gq = cp.quantile(norms, levels).astype(cp.float32) # global fallback grid + cid = cp.asarray(cluster_ids) + q = cp.empty((nc, _NORM_QUANTILE_COUNT), dtype=cp.float32) + for c in range(nc): + m = cid == c + q[c] = ( + cp.quantile(norms[m], levels).astype(cp.float32) + if bool(m.any()) + else gq + ) + cp.get_default_memory_pool().free_all_blocks() + return cp.asnumpy(q), mean, cv + + +def sample_norms_percentile(cluster_ids, norm_quantiles, seed=42): + """Draw each node's target norm from its cluster's inverse-CDF (percentile + scheme). Mirrors synthesize_dataset/_generate.py `_rescale_to_scheme`. + """ + rng = np.random.default_rng(seed + 4) + nq = np.ascontiguousarray(norm_quantiles, dtype=np.float32) # (nc, Q) + Q = nq.shape[1] + cid = cluster_ids.astype(np.int64) + u = rng.random(cid.shape[0]).astype(np.float32) # (N,) + pos = u * (Q - 1) + lo = np.minimum(np.floor(pos).astype(np.int64), Q - 2) # left grid index + frac = (pos - lo).astype(np.float32) + q_lo = nq[cid, lo] # (N,) gather + q_hi = nq[cid, lo + 1] + return (q_lo * (1.0 - frac) + q_hi * frac).astype(np.float32) + + +@torch.no_grad() +def decode_graph( + anchor_ids, + nbr, + model, + anchors, + mu, + sd, + feat_mu, + feat_sd, + resid_params, + norm_target, + device, + seed, + resid_ids=None, + base_path=None, + query_idx=None, +): + """Decode a kNN graph, streaming base rows to disk -> return held-out queries (N_q, D). + Runs the trained decoder over the N-node graph to get the vector, adds the + per-cluster residual, and rescales each vector's radius. Base rows are written + straight to base_path and only the query rows are kept and returned. + + Parameters + ---------- + anchor_ids : (N,) int32 anchor id per node — indexes the decoder's anchor + embedding table. + nbr : (N, k) int64 the generated (synthetic) kNN graph to decode. + model : the trained kNNDecoder. + anchors : (n_anchors, D) float32 the sample rows X (host); nodes index into + this table by anchor id to get their anchor embedding. + mu, sd : (1, D) float32 per-dim stats of the real sample; used to + un-normalize the decoder output back into ambient space. + feat_mu, : structural-feature normalization stats from train_model, so the + feat_sd synthetic graph's features are scaled exactly as in training. + resid_params : list of per-cluster residual Gaussians from fit_residuals + norm_target : (N,) float32 target L2 norm per node (drawn from the per-cluster + norm inverse-CDF) + device : torch device for the decode. + seed : RNG seed for the residual sampling. + resid_ids : (N,) int32 the KMeans cluster label per node + base_path : STREAM base rows straight to this .fbin as we decode to cap host RAM usage. + query_idx : (N_q,) sorted int64 ids of the held-out query rows + """ + N, k = nbr.shape + D = mu.shape[1] + chunk = CHUNK_SIZE + if resid_ids is None: + resid_ids = anchor_ids + torch.manual_seed(seed + 1) + mu_t = torch.tensor(mu, device=device) + sd_t = torch.tensor(sd, device=device) + struct_all = compute_structural_features( + nbr + ) # in-degrees of the synthetic graph + struct_all = normalize_features(struct_all, feat_mu, feat_sd) + + # Precompute the anchor-embedding table ONCE + n_anc = anchors.shape[0] + d_emb = model.cluster_mlp[-1].out_features + emb_table = torch.empty(n_anc, d_emb, device=device) + for i in range(0, n_anc, chunk): + end = min(i + chunk, n_anc) + emb_table[i:end] = model.cluster_mlp( + torch.as_tensor(anchors[i:end], device=device) + ) + + # Base rows stream straight to disk; only the held-out queries stay in RAM. + n_q = len(query_idx) + queries = np.empty((n_q, D), dtype=np.float32) + q_off = 0 + fbase = open(base_path, "wb") + write_fbin_header(fbase, N - n_q, D) # base row count known up front + + # Overlap disk writes with the GPU decode: + write_thread = None + write_exception = None + + def _wait_write(): + nonlocal write_thread, write_exception + if write_thread is not None: + write_thread.join() + write_thread = None + if write_exception is not None: + exc = write_exception + write_exception = None + raise exc + + def _flush_async(arr): + nonlocal write_thread + + def _w(): + nonlocal write_exception + try: + arr.tofile(fbase) + except BaseException as e: + write_exception = e + + write_thread = threading.Thread(target=_w, daemon=True) + write_thread.start() + + for start in tqdm(range(0, N, chunk), desc="decode", unit="chunk"): + end = min(start + chunk, N) + b = end - start + + # --- center-node features: anchor embedding (gathered) + structural features --- + anchor_i = torch.tensor( + anchor_ids[start:end].astype(np.int64), device=device + ) + struct_i = torch.tensor(struct_all[start:end], device=device) + hi = torch.cat([emb_table[anchor_i], struct_i], dim=-1) # (b, d_node) + + # --- neighbor features: same, for every node's k neighbours --- + flat = nbr[start:end].reshape(-1) # (b*k,) flattened neighbour ids + anchor_j = torch.tensor( + anchor_ids[flat].astype(np.int64), device=device + ) + struct_j = torch.tensor(struct_all[flat], device=device) + hj = torch.cat([emb_table[anchor_j], struct_j], dim=-1).reshape( + b, k, -1 + ) # (b, k, d_node) + + # --- decode features into vectors --- + pred = model(hi, hj) # (b, D) in z-normalized space + + # --- add the per-cluster residual (restores within-cluster spread) --- + resid_i = torch.tensor( + resid_ids[start:end].astype(np.int64), device=device + ) + pred = pred + sample_residuals(resid_i, resid_params, D, device) + + # --- un-normalize back to ambient space --- + xb = pred * sd_t + mu_t + + # --- radial rescale: keep only the direction, set the L2 norm to the target --- + dirv = xb / (xb.norm(dim=1, keepdim=True) + 1e-12) # unit direction + r = torch.tensor(norm_target[start:end], device=device).unsqueeze( + 1 + ) # target radius + xb = dirv * r # vector = direction * target norm + + xb_np = xb.cpu().numpy() # (b, D) fresh host tile per chunk + # which rows in [start, end) are held-out queries — sliced from the sorted + # id list via searchsorted (no O(N) mask). + lo, hi = ( + np.searchsorted(query_idx, start), + np.searchsorted(query_idx, end), + ) + n_qc = hi - lo + if n_qc: + qm = np.zeros(b, dtype=bool) + qm[query_idx[lo:hi] - start] = True + queries[q_off : q_off + n_qc] = xb_np[ + qm + ] # keep the few query rows in RAM + q_off += n_qc + base = np.ascontiguousarray( + xb_np[~qm] + ) # base rows the writer thread owns + else: + base = xb_np # no queries here: the whole tile is base + + _wait_write() # prev flush done -> its buffer is free + _flush_async(base) # write this chunk while the next decodes + + _wait_write() + fbase.close() + return queries diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/upsample.py b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/upsample.py new file mode 100644 index 0000000000..4d85e83f14 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/upsample.py @@ -0,0 +1,196 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""kNN upsampling: build the N-node synthetic graph from the sample kNN.""" + +from __future__ import annotations + +import numpy as np +from tqdm import tqdm +from utils import CHUNK_SIZE + + +def _node_coherence(nbr, idx): + """Triangle Density for the nodes in `idx`. Returns (len(idx),) float32: fraction + of each node's ordered neighbor-pairs (j, j') where j' is also its neighbor. + """ + k = nbr.shape[1] + nb = nbr[idx] # (m, k) i.e. the k nbrs + nnb = nbr[nb] # (m, k, k) nbrs-of-nbrs (2-hop) + match = (nnb[..., None] == nb[:, None, None, :]).any( + -1 + ) # (m, k, k) j' in N(i)? + return (match.sum((1, 2)) / (k * (k - 1))).astype( + np.float32 + ) # (m) frac btw [0,1] + + +def _est_coherence(nbr, seed=0, m=4000): + """Mean local coherence over a random sample of m nodes + — a cheap scalar for the triadic window-sizing loop. + """ + n_nodes, k = nbr.shape + if k < 2: + return 0.0 + rng = np.random.default_rng(seed) + idx = rng.integers(0, n_nodes, size=min(m, n_nodes)) + return float(_node_coherence(nbr, idx).mean()) + + +def blend_chunglu(core_nbr, k_keep, indeg_dist, N, k, seed, chunk=CHUNK_SIZE): + """Blend a coherent/structured neighbor matrix with a Chung-Lu hub backbone. + + For each node i, keep its first ``k_keep[i]`` core neighbors and fill the rest of + its k edges with global degree-weighted random draws: node attractiveness is + sampled from the real sample in-degree distribution (indeg_dist) and targets are drawn + proportional to it, so a few nodes accumulate many in-edges — the heavy + in-degree tail (hubs) the index leans on. Returns (N, k) int64. + + Note: the result may contain DUPLICATE neighbors and self-loops. This is not a problem because + `nbr` is a throwaway decoder input. A duplicate just gives that neighbor a little extra + weight in the order-invariant attention aggregate. + """ + import cupy as cp + + rng = np.random.default_rng(seed) + weights = ( + rng.choice(indeg_dist, size=N).astype(np.float64) + 1e-6 + ) # per-node attractiveness + cdf = cp.cumsum(cp.asarray(weights)) # inverse-CDF over N targets + cdf /= cdf[-1] # -> [0, 1] + del weights + + out = np.empty((N, k), dtype=np.int64) + ar_k = np.arange(k) + crng = cp.random.RandomState(seed + 11) + for s in tqdm(range(0, N, chunk), desc="chung-lu blend", unit="chunk"): + e = min(s + chunk, N) + b = e - s + u = crng.random_sample((b * k,)) # uniform draws + rand = cp.searchsorted(cdf, u, side="right").reshape( + b, k + ) # (b, k) targets ∝ weight + rand = cp.asnumpy(cp.minimum(rand, N - 1)).astype(np.int64) + mask = ar_k[None, :] < k_keep[s:e, None] # first k_keep are core + out[s:e] = np.where( + mask, core_nbr[s:e], rand + ) # dups/self-loops possible; see note + del cdf + cp.get_default_memory_pool().free_all_blocks() + return out + + +def triadic_coherent_knn(cluster_ids, k, seed, tgt): + """Coordinate-free coherent graph. + + Within each cluster, lay the nodes in a random order and connect each to k random + nodes inside a sliding window of width w. Overlapping windows produce continuous, + mutually-overlapping neighborhoods (triangles). The window width sets the coherence, + so w is auto-sized to the sample's measured coherence (tgt) with one proportional + correction. A smaller window results in a higher coherence. + + Returns (N, k) int64 + """ + import cupy as cp + + N = cluster_ids.shape[0] + + order = np.argsort( + cluster_ids, kind="stable" + ) # order[p] = actual global node id at position p + starts, counts = np.unique( + cluster_ids[order], return_index=True, return_counts=True + )[1:] + bstart = np.repeat(starts, counts).astype( + np.int64 + ) # (N,) bstart[p] = where p's cluster block starts + sz_pos = np.repeat(counts, counts).astype( + np.int64 + ) # (N,) sz_pos[p] = p's cluster size + loc = (np.arange(N) - bstart).astype( + np.int64 + ) # (N,) loc[p] = p-bstart[p] = p's local index inside cluster + order_d = cp.asarray(order.astype(np.int64)) # global ids, on GPU + + def build(w, bseed, chunk=CHUNK_SIZE): + # Each node connects to k random nodes within +/- w/2 positions inside its + # cluster block (overlapping windows -> triangles). Done in node-chunks on + # the GPU. + half = max(1, w // 2) + crng = cp.random.RandomState(bseed) + nbr = np.empty((N, k), dtype=np.int64) + for s in tqdm( + range(0, N, chunk), desc=f"wiring coherence (w={w})", unit="chunk" + ): + e = min(s + chunk, N) + loc_c = cp.asarray(loc[s:e])[:, None] # (b,1) + sz_c = cp.asarray(sz_pos[s:e])[:, None] # (b,1) + bst_c = cp.asarray(bstart[s:e])[:, None] # (b,1) + off = crng.randint( + -half, half + 1, size=(e - s, k) + ) # (b,k) window offsets + tp = (loc_c + off) % sz_c # local target positions + tp = cp.where(tp == loc_c, (tp + 1) % sz_c, tp) # avoid self + tgt = order_d[bst_c + tp] # (b,k) local -> global + nbr[order[s:e]] = cp.asnumpy(tgt) # scatter to node order + return nbr + + w = max(k + 1, round(0.7 * k / max(tgt, 1e-3))) + nbr = build(w, seed + 7) + C = _est_coherence(nbr, seed=1) + if abs(C - tgt) > 0.015 and C > 0: # one proportional correction (C ~ 1/w) + w = max(k + 1, round(w * C / tgt)) + nbr = build(w, seed + 8) + C = _est_coherence(nbr, seed=1) + print( + f" coherent seed: windowed in-cluster (w={w}), " + f"coherence~{C:.3f} (target~{round(tgt, 3)})", + flush=True, + ) + return nbr + + +def generate_graph_knn(stats, N, k, seed, knn_frac=1.0): + """Build coherent N-node kNN graph, blended with a Chung-Lu hub tail. + + Generate a coherent kNN graph using triadic_coherent_knn. blend_chunglu then + keeps ~knn_frac*k coherent edges per node and fills the rest with degree-weighted + random (Chung-Lu) edges. The coherentedges drive recall, the random tail drives + the hub tail + search difficulty (knn_frac=1.0 = fully coherent). Larger knn_frac + means an easier data. + + Returns (anchor_ids (N,) int32 — each node's anchor, nbr (N, k) int64). + """ + import cupy as cp + + rng = np.random.default_rng(seed) + indeg_dist = stats["indeg_dist"] + n_anchors = len(indeg_dist) # #anchors == sample size + + anchor_ids = rng.integers(0, n_anchors, size=N).astype(np.int32) + knn_nbr = triadic_coherent_knn( + anchor_ids, k, seed=seed, tgt=stats.get("coherence") + ) + cp.get_default_memory_pool().free_all_blocks() + + kf_k = knn_frac * k + if int(np.floor(kf_k)) >= k: + return anchor_ids, knn_nbr + + # Per-node coherent count so the effective fraction is continuous + base = int(np.floor(kf_k)) + fracp = kf_k - base + k_knn_i = np.clip( + base + (rng.random(N) < fracp).astype(np.int64), 0, k + ) # (N,) + + # Blend: keep each node's k_knn_i coherent kNN neighbors, fill the rest with + # the shared Chung-Lu hub backbone. + nbr = blend_chunglu(knn_nbr, k_knn_i, indeg_dist, N, k, seed + 5) + print( + f" blend: ~{kf_k:.2f} coherent kNN + ~{k - kf_k:.2f} chung-lu random " + f"per node (knn_frac={knn_frac}, per-node stochastic)", + flush=True, + ) + return anchor_ids, nbr diff --git a/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/utils.py b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/utils.py new file mode 100644 index 0000000000..20762e60e9 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/synthesize_dataset/knn_decode_synthesizer/utils.py @@ -0,0 +1,313 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""Shared helpers: I/O, logging, cuVS wrappers, and the per-cluster residual sampler.""" + +from __future__ import annotations + +import os +import struct +import time + +import numpy as np +import torch +from tqdm import tqdm + +# Row-chunk for every streaming GPU pass (decode, residual fit, norms, mean/std, +# graph wiring, brute-force GT tiles). Caps peak memory to ~one chunk's worth of +# rows. +CHUNK_SIZE = 1_000_000 + + +def load_fbin(path): + """Load an .fbin file (``[n, d]`` int32 header + ``n*d`` float32) as (n, d).""" + with open(path, "rb") as f: + n, d = struct.unpack(" (B, D) float32 tensor (zeros for clusters with no fitted Gaussian) + """ + B = cluster_chunk.shape[0] + out = torch.zeros(B, D, device=device) + for c in torch.unique(cluster_chunk).tolist(): + p = resid_params[c] + if p is None: + continue + sel = (cluster_chunk == c).nonzero(as_tuple=True)[0] + n = sel.shape[0] + z = torch.randn(n, p["comps"].shape[0], device=device) * p["stds"] + s = z @ p["comps"] # (n, D) + eps = torch.randn(n, D, device=device) * p["noise_std"] + out[sel] = p["mean"] + s + eps + return out + + +def write_fbin(path, A): + """Write a 2-D array to a cuvs-bench .fbin (uint32 [n, d] header + float32 data).""" + A = np.ascontiguousarray(A, dtype=np.float32) + with open(path, "wb") as f: + f.write(struct.pack(" global row id + merged_i = np.concatenate([gt, ii], axis=1) + merged_d = np.concatenate([gtd, dd], axis=1) + order = np.argsort(merged_d, axis=1)[:, :k] + gt = np.take_along_axis(merged_i, order, axis=1) + gtd = np.take_along_axis(merged_d, order, axis=1) + del tile_d, idx + with open(os.path.join(out_dir, "groundtruth.neighbors.ibin"), "wb") as f: + f.write(struct.pack(" {out_dir}", + flush=True, + ) + _write_exact_gt( + out_dir, + queries, + len(base), + gt_k, + lambda s, e: np.ascontiguousarray(base[s:e], dtype=np.float32), + tile, + ) + + +def write_bundle_streamed( + out_dir, base_path, base_count, queries, gt_k, d, tile=CHUNK_SIZE +): + """Finalize a bundle whose base.fbin was already STREAMED to disk by the decoder + (see decode_graph's streaming mode). Writes queries.fbin and exact brute-force + GT, reading the base back from disk one tile at a time — so the full base never + has to be resident in host RAM nor on the GPU. + """ + write_fbin(os.path.join(out_dir, "queries.fbin"), queries) + print( + f" base {base_count:,} x {d} (streamed to disk) " + f"queries {queries.shape} -> {out_dir}", + flush=True, + ) + _write_exact_gt( + out_dir, + queries, + base_count, + gt_k, + lambda s, e: read_fbin_rows(base_path, s, e, d), + tile, + )