Skip to content

Add HNSW Layered Index Support#2148

Draft
julianmi wants to merge 7 commits into
rapidsai:mainfrom
julianmi:hnsw-layered-index
Draft

Add HNSW Layered Index Support#2148
julianmi wants to merge 7 commits into
rapidsai:mainfrom
julianmi:hnsw-layered-index

Conversation

@julianmi
Copy link
Copy Markdown
Contributor

@julianmi julianmi commented Jun 1, 2026

The CAGRA graph built by the disk-backed ACE algorithm partitions the dataset. Thus, the CAGRA graph uses the reordered index space. Building a HNSW index using hnsw::from_cagra uses the reordered dataset and CAGRA graph. Downstream consumers building an HNSW index would therefore require the reordered dataset, which is typically large when requiring the disk-backed ACE algorithm. Thus, building only the layers of the HNSW index without the dataset and moving this to the search node can minimize the network transfers for downstream consumers if they have the original dataset locally available. The hnsw::deserialize step then takes the layered index and combines it with the local dataset to form a hnswlib compatible search index.

Artifact Layout

hnsw_index.cuvs
  fixed_header
    magic = CUVS_HNSW_LAYERED
    version
    metadata_offset
    metadata_size

  metadata_json
    dataset shape, dtype, metric
    hnsw parameters
    section sizes
    upper-layer descriptors

  levels
    uint8[n_rows]
    indexed by original dataset row ID

  base_nodes
    uint32[n_rows]
    maps each base topology row to original row ID

  base_links
    n_rows fixed-size hnswlib-ready rows
    [count:uint32][neighbors:uint32[maxM0]]
    neighbors are original row IDs

  upper_nodes
    concatenated uint32 original row IDs for layers 1..maxlevel

  upper_links
    fixed-size hnswlib-ready rows
    [count:uint32][neighbors:uint32[maxM]]
    neighbors are original row IDs

Layered HNSW Serialization

The layered serializer creates hnsw_index.cuvs from the disk-backed ACE graph.

  1. Create the .cuvs file and write the fixed header and metadata JSON.
  2. Generate HNSW levels in original ID space.
  3. Write levels sequentially.
  4. Read dataset_mapping.npy sequentially into reordered_to_original.
  5. Read cagra_graph.npy source-sequentially in ACE reordered row order.
  6. For each ACE graph row:
    • write base_nodes[row] = reordered_to_original[ace_reordered_row]
    • convert each neighbor from ACE reordered ID to original ID
    • write a padded hnswlib-ready row to base_links[row]
  7. Gather promoted vectors from the original dataset.
  8. Build upper-layer graphs using temporary HNSW promoted order.
  9. Write upper_nodes and upper_links with node IDs and neighbor IDs converted back to original IDs.

This keeps remapping, link padding, and upper-layer KNN work on the build node.

Deserialization

The search node reads:

  • hnsw_index.cuvs
  • the external original-order dataset from index_params.dataset_path

The loader:

  1. Reads the fixed header and metadata.
  2. Validates artifact shape, section sizes, and dataset shape.
  3. Reads levels sequentially.
  4. Allocates hnswlib storage.
  5. Reads the external dataset sequentially in original row order.
  6. Initializes hnswlib with:
    • internal ID = original ID
    • label = original ID
    • level = levels[original_id]
  7. Reads base_nodes and base_links sequentially.
  8. Copies each base link row into get_linklist0(base_node_id).
  9. Reads upper_nodes and upper_links sequentially by layer.
  10. Copies each upper link row into get_linklist(node_id, level).

The search node does no graph remapping, no level generation, no link padding, and no KNN work.

Disk Access Patterns

Build node:

  • Sequential scan of the original dataset for ACE partitioning.
  • Buffered partition writes for reordered and augmented datasets.
  • Contiguous per-partition reads from reordered_dataset.npy and augmented_dataset.npy.
  • Source-sequential reads from cagra_graph.npy when creating the final layered artifact.
  • Sequential writes to hnsw_index.cuvs.

Search node:

  • Sequential reads from hnsw_index.cuvs.
  • Sequential reads from the external original-order dataset.
  • Scatter writes only into in-memory hnswlib link storage by original ID.

Runtime Requirements

Only hnsw_index.cuvs is copied to the search node. ACE temporary files remain build-node-only.

The search node must have the original dataset in original row order and must provide that path through index_params.dataset_path.

Misc

Unifies the logging format of the ACE algorithm.

@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented Jun 1, 2026

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@tfeher tfeher requested a review from mfoerste4 June 1, 2026 09:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant