Binary dataset support in CAGRA (upstream #1906)

[mohanprasand-nuvai]

Context

Upstream rapidsai/cuvs#1906 requests binary dataset support in CAGRA — specifically, the ability to accept binary-quantized vectors with residual codes for ADC (asymmetric distance computation). This is a requirement from OpenSearch.

Upstream Dependency

This feature is blocked upstream on rapidsai/cuvs#1846 — a WIP PR that introduces a new padded_dataset API replacing strided_dataset, with support plumbed into CAGRA. Key design decisions in NVIDIA#1846:

1. Users should not need to know about dataset padding internals — padded_dataset handles it
2. Graph-based APIs (CAGRA, Vamana) should accept pre-quantized data (e.g., PQ with codebooks) rather than performing their own quantization — decouples quantization from the algorithm
3. Memory ownership stays with the user, not the algorithm

Status (as of 2026-03-23): PR NVIDIA#1846 is still open and marked WIP. Last activity: March 4, 2026.

Decision Needed

Two options:

Option A: Cherry-pick NVIDIA#1846 + implement NVIDIA#1906

• Wait for [REVIEW] New Dataset API Clarifying Ownership NVIDIA/cuvs#1846 to stabilize, cherry-pick it, then build [FEA] Implement binary dataset and accept in CAGRA NVIDIA/cuvs#1906 on top
• Pro: Aligned with upstream direction, easier to merge future upstream changes
• Con: [REVIEW] New Dataset API Clarifying Ownership NVIDIA/cuvs#1846 is a breaking change (strided_dataset → padded_dataset), large surface area, still WIP

Option B: From-scratch implementation

• Implement binary dataset support directly without the new dataset API
• Could add a binary_dataset class that wraps packed bit vectors and provides Hamming distance
• Wire it into CAGRA build/search via if constexpr dispatch (same pattern used for double/bf16 in PR feat: add float64 and bfloat16 support to CAGRA #4)
• Pro: No dependency on upstream WIP, ships sooner, smaller scope
• Con: May diverge from upstream's eventual API, potential rework when [REVIEW] New Dataset API Clarifying Ownership NVIDIA/cuvs#1846 merges

Option C: Hybrid — minimal dataset abstraction now, align later

• Implement a lightweight binary_dataset that's compatible with the direction of [REVIEW] New Dataset API Clarifying Ownership NVIDIA/cuvs#1846 but doesn't require the full padded_dataset refactor
• Pro: Ships sooner, easier to align when [REVIEW] New Dataset API Clarifying Ownership NVIDIA/cuvs#1846 merges
• Con: Requires understanding [REVIEW] New Dataset API Clarifying Ownership NVIDIA/cuvs#1846's design intent well enough to not paint ourselves into a corner

Scope (for any option)

• Binary dataset type (packed bit vectors, likely uint8_t storage with N vectors of D bits → [N × ceil(D/8)] uint8 matrix)
• Hamming distance computation for CAGRA search kernels
• CAGRA build with binary data (likely using float-promoted intermediate graph, similar to int8/bf16 pattern)
• C-API dispatch for binary datasets
• Rust/Go bindings

References

• Upstream issue: [FEA] Implement binary dataset and accept in CAGRA NVIDIA/cuvs#1906
• Blocking PR: [REVIEW] New Dataset API Clarifying Ownership NVIDIA/cuvs#1846
• Related upstream issues: [FEA] Standardize Datasets API for public end-user use NVIDIA/cuvs#1574, [FEA] Indexes to accept Dataset directly NVIDIA/cuvs#1571 (dataset ownership/padding)

[mohanprasand-nuvai]

Implemented in PR #7 — Option C (Hybrid)

Fresh implementation following the existing vpq_dataset pattern. No dependency on upstream PR NVIDIA#1846.

What shipped

Part 1: binary_dataset class (common.hpp)

• binary_dataset<IdxT> inheriting dataset<IdxT>, packed uint8_t storage [N, ceil(D/8)]
• dim() returns original bit dimensionality, packed_dim() returns byte width
• Constructor validates data.extent(1) >= ceil(dim/8)
• is_binary_dataset_v<T> type trait

Part 2: ADC search kernel (new files)

• binary_adc_descriptor_spec — registered in code generator, 3 instantiations (TeamSize 8/16/32)
• binary_adc_dataset_descriptor_t — branchless float(bit) * query_val per dimension
• Per-byte ld.global.cg.u8 loads (alignment-safe for any packed_dim)
• Distance: -dot(query, binary_vec) — correct ranking proxy for L2 and InnerProduct
• Metric validation: only L2Expanded and InnerProduct (L1 semantically wrong, L2Sqrt→NaN, Cosine needs norms)

Part 3: Search integration (cagra_search.cuh, factory.cuh)

• dynamic_cast<binary_dataset> branch in search_main
• RAFT_EXPECTS guards for unsupported metrics
• make_key overload for descriptor cache

Part 4: Go binding fix (distance.go)

• Added DistanceBitwiseHamming constant + C.CUVS_DISTANCE_BITWISE_HAMMING mapping

Part 5: Rust binding (distance_type.rs)

• Added pub const BITWISE_HAMMING: DistanceType re-export

Part 6: Tests

• AnnCagraBinaryDatasetTest — unit test for dim/packed_dim/data integrity/validation
• AnnCagraBinaryAdcTest — E2E: float build → binary quantize → ADC search → recall check (clustered data)
• Go: TestDistanceBitwiseHammingConstant
• Rust: test_bitwise_hamming_equals_ffi + test_bitwise_hamming_value

What's pending (fast follows, not blocking)

• Part 3 from plan: Oversampling + re-ranking utility (search_with_reranking) — for >90% recall by oversampling with ADC then re-ranking against float originals
• Serialization support — binary_dataset is not persisted across cagra::serialize/deserialize; user must re-attach after deserialize
• C-API dispatch — binary dataset workflow is currently C++ only; no cuvsCagra* C function changes
• E2E symmetric Hamming test (Plan Part 6.1) — validates uint8_t build + BitwiseHamming + NN-Descent path (existing infra already tests this parametrically, but a dedicated focused test was planned)