From 3f1d7228b24eb9c304b4379745fccb6a2cba753e Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Thu, 1 Oct 2026 12:53:47 +0000 Subject: [PATCH] Fix scalar CAGRA graph byte encoding --- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 2 +- ...leratedhnswscalarquantizedvectorswriter.md | 10 +- .../cuvs/lucene/AcceleratedHNSWUtils.java | 10 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 31 +-- .../cuvs/lucene/TestScalarQuantization.java | 41 ++++ .../lucene/TestScalarQuantizedCagraGraph.java | 190 ++++++++++++++++++ 6 files changed, 248 insertions(+), 36 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 219149bc6d..22cb8786e9 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -160,7 +160,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWU public static List quantizeFloatVectorsToScalar(List floatVectors) ``` -Scalar quantization. +Scalar quantization to unsigned 7-bit values stored in Java bytes. **Parameters** diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 9c51cb26e9..e3e66b1e45 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:241`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:222`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:326`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:307`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:357`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:338`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:377`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:358`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:386`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:367`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 9c49c07fe0..27ff26e460 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -444,7 +444,7 @@ public static List quantizeFloatVectorsToBinary(List floatVecto } /** - * Scalar quantization. + * Scalar quantization to unsigned 7-bit values stored in Java bytes. * * @param floatVectors A list of float vectors * @return A list of byte scalar representation for the input vectors @@ -459,8 +459,8 @@ public static List quantizeFloatVectorsToScalar(List floatVecto float[] minPerDim = new float[dimensions]; float[] maxPerDim = new float[dimensions]; - Arrays.fill(minPerDim, Float.MAX_VALUE); - Arrays.fill(maxPerDim, Float.MIN_VALUE); + Arrays.fill(minPerDim, Float.POSITIVE_INFINITY); + Arrays.fill(maxPerDim, Float.NEGATIVE_INFINITY); for (float[] vector : floatVectors) { for (int d = 0; d < dimensions; d++) { @@ -476,8 +476,8 @@ public static List quantizeFloatVectorsToScalar(List floatVecto float range = maxPerDim[d] - minPerDim[d]; if (range > 0) { float normalized = (vector[d] - minPerDim[d]) / range; - int quantizedValue = Math.round(normalized * 127.0f) - 64; - quantized[d] = (byte) Math.max(-64, Math.min(63, quantizedValue)); + int quantizedValue = Math.round(normalized * 127.0f); + quantized[d] = (byte) Math.max(0, Math.min(127, quantizedValue)); } else { quantized[d] = 0; } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 7141af56ee..75b97b5d97 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -146,18 +146,6 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException return writer; } - private static byte signedToUnsignedByte(byte signedByte) { - return (byte) (signedByte & 0xFF); - } - - private static byte[] convertSignedToUnsigned(byte[] signedVector) { - byte[] unsignedVector = new byte[signedVector.length]; - for (int i = 0; i < signedVector.length; i++) { - unsignedVector[i] = signedToUnsignedByte(signedVector[i]); - } - return unsignedVector; - } - /** * Builds the intermediate CAGRA index and builds and writes the HNSW index. * @@ -165,7 +153,7 @@ private static byte[] convertSignedToUnsigned(byte[] signedVector) { * @param vectors quantized vectors * @throws IOException */ - private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { if (vectors.size() == 0) { writeEmpty(fieldInfo, hnswMeta); return; @@ -174,18 +162,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE try { int dimensions = fieldInfo.getVectorDimension(); - // Convert 7-bit signed bytes to 8-bit unsigned bytes for cuVS compatibility - List unsignedVectors = new ArrayList<>(vectors.size()); - for (Object signedVector : vectors) { - unsignedVectors.add(convertSignedToUnsigned((byte[]) signedVector)); - } - - // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = - Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); + // The scalar quantizer already emits nonnegative bytes for cuVS's unsigned BYTE type. + CuVSMatrix dataset = Utils.createByteMatrix(vectors, dimensions, getCuVSResourcesInstance()); if (dataset.size() < 2) { - writeSingleVectorGraph(fieldInfo, unsignedVectors); + writeSingleVectorGraph(fieldInfo, vectors); return; } @@ -207,7 +188,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE size, dimensions, adjacencyListMatrix, - unsignedVectors, + vectors, acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.SCALAR); @@ -287,7 +268,7 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo - * @param vectors the list of scalar quantized vectors (already converted to unsigned) + * @param vectors scalar-quantized nonnegative byte vectors * @throws IOException I/O Exceptions */ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java new file mode 100644 index 0000000000..48378e865b --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.junit.Assert.assertArrayEquals; + +import java.util.List; +import org.junit.Test; + +public class TestScalarQuantization { + + @Test + public void quantizesNegativeOnlyDimensionsAcrossTheFullRange() { + List quantized = + AcceleratedHNSWUtils.quantizeFloatVectorsToScalar( + List.of(new float[] {-10.0f, -4.0f}, new float[] {-5.0f, -2.0f})); + + assertArrayEquals(new byte[] {0, 0}, quantized.get(0)); + assertArrayEquals(new byte[] {127, 127}, quantized.get(1)); + } + + @Test + public void quantizesIntermediateValuesMonotonically() { + List quantized = + AcceleratedHNSWUtils.quantizeFloatVectorsToScalar( + List.of( + new float[] {-1.0f}, + new float[] {-0.5f}, + new float[] {0.0f}, + new float[] {0.5f}, + new float[] {1.0f})); + + assertArrayEquals(new byte[] {0}, quantized.get(0)); + assertArrayEquals(new byte[] {32}, quantized.get(1)); + assertArrayEquals(new byte[] {64}, quantized.get(2)); + assertArrayEquals(new byte[] {95}, quantized.get(3)); + assertArrayEquals(new byte[] {127}, quantized.get(4)); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java new file mode 100644 index 0000000000..8e4a8357f6 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java @@ -0,0 +1,190 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.assertIsSupported; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.junit.Assume.assumeTrue; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.stream.IntStream; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.junit.Test; + +/** Checks the GPU-built scalar graph with CPU HNSW search against exact Euclidean neighbors. */ +public class TestScalarQuantizedCagraGraph extends LuceneTestCase { + private static final String VECTOR_FIELD = "vector"; + private static final String ID_FIELD = "id"; + private static final int VECTOR_COUNT = 512; + private static final int DIMENSIONS = 128; + private static final int TOP_K = 10; + private static final int MIN_EXACT_NEIGHBORS = 8; + + @Test + public void testGpuBuiltScalarHnswRetainsRecallAcrossMerge() throws Exception { + requireGpuWhenSelected(); + + float[][] vectors = vectorsWithNegativeAndMixedSignDimensions(); + RecordingInfoStream buildLog = new RecordingInfoStream(); + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .withHNSWLayer(1) + .build(); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec( + TestUtil.alwaysKnnVectorsFormat( + new LuceneAcceleratedHNSWScalarQuantizedVectorsFormat(params))) + .setMaxBufferedDocs(VECTOR_COUNT + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setInfoStream(buildLog); + + try (Directory directory = newDirectory()) { + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int id = 0; id < vectors.length; id++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(id), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[id], EUCLIDEAN)); + writer.addDocument(document); + if (id == VECTOR_COUNT / 2 - 1) { + writer.commit(); + } + } + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(2, reader.leaves().size()); + assertHnswRecallAgainstExactNeighbors(reader, vectors); + } + + long gpuBuildsBeforeMerge = buildLog.gpuWriterOpenCount(); + assertTrue("No scalar GPU writer opened: " + buildLog.messages, gpuBuildsBeforeMerge >= 2); + + writer.forceMerge(1); + writer.commit(); + assertTrue( + "The merge did not open a scalar GPU writer: " + buildLog.messages, + buildLog.gpuWriterOpenCount() > gpuBuildsBeforeMerge); + } + + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(1, reader.leaves().size()); + assertHnswRecallAgainstExactNeighbors(reader, vectors); + } + } + } + + private static void requireGpuWhenSelected() { + if (Boolean.getBoolean("cuvs.lucene.tests.requireGpu")) { + assertIsSupported(); + } else { + assumeTrue("cuVS is not supported", isSupported()); + } + } + + private static void assertHnswRecallAgainstExactNeighbors( + DirectoryReader reader, float[][] vectors) throws Exception { + IndexSearcher searcher = new IndexSearcher(reader); + for (int queryId : new int[] {0, 240, 272, VECTOR_COUNT - 1}) { + var results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, vectors[queryId], TOP_K), TOP_K); + List actual = new ArrayList<>(); + for (var hit : results.scoreDocs) { + actual.add(Integer.parseInt(searcher.storedFields().document(hit.doc).get(ID_FIELD))); + } + + assertEquals("Query " + queryId, TOP_K, actual.size()); + assertEquals("Query " + queryId, queryId, actual.get(0).intValue()); + assertEquals( + "Query " + queryId + " returned duplicates", TOP_K, new HashSet<>(actual).size()); + + List exact = + IntStream.range(0, vectors.length) + .boxed() + .sorted( + Comparator.comparingDouble( + (Integer id) -> squaredDistance(vectors[queryId], vectors[id])) + .thenComparingInt(Integer::intValue)) + .limit(TOP_K) + .toList(); + Set expected = new HashSet<>(exact); + long overlap = actual.stream().filter(expected::contains).count(); + assertTrue( + "Query " + queryId + " found " + overlap + "/" + TOP_K + " exact neighbors: " + actual, + overlap >= MIN_EXACT_NEIGHBORS); + } + } + + private static double squaredDistance(float[] left, float[] right) { + double distance = 0; + for (int dimension = 0; dimension < left.length; dimension++) { + double difference = left[dimension] - right[dimension]; + distance += difference * difference; + } + return distance; + } + + private static float[][] vectorsWithNegativeAndMixedSignDimensions() { + float[][] vectors = new float[VECTOR_COUNT][DIMENSIONS]; + for (int id = 0; id < VECTOR_COUNT; id++) { + int x = id % 32; + int y = id / 32; + for (int dimension = 0; dimension < DIMENSIONS; dimension++) { + vectors[id][dimension] = + switch (dimension % 4) { + case 0 -> -20.0f + 0.4f * x + 0.006f * y; // all negative + case 1 -> -7.5f + 0.9f * y + 0.002f * x; // crosses zero + case 2 -> 4.0f + 0.25f * x + 0.003f * y; // all positive + default -> -6.0f + 0.21f * (x + y); // crosses zero + }; + } + } + return vectors; + } + + private static final class RecordingInfoStream extends InfoStream { + private static final String GPU_WRITER_OPENED = + "Lucene99AcceleratedHNSWQuantizedVectorsWriter opened"; + private final List messages = new CopyOnWriteArrayList<>(); + + @Override + public void message(String component, String message) { + messages.add(component + ": " + message); + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + + private long gpuWriterOpenCount() { + return messages.stream().filter(message -> message.contains(GPU_WRITER_OPENED)).count(); + } + } +}