From c87fcd47e0a4319840d6af6c63068abba3d17248 Mon Sep 17 00:00:00 2001 From: Igor Motov Date: Tue, 18 Aug 2026 10:51:19 -1000 Subject: [PATCH] Fix silent brute force fallback at aligned vector dimensions cuvsDatasetMakePadded rejects a device matrix already at CAGRA's required stride, so every field whose dimension is a multiple of four failed its CAGRA build and was silently indexed as brute force. Pick the factory that matches the layout, and log the fallback so the next one is visible. See #2482 --- .../main/java/com/nvidia/cuvs/CagraIndex.java | 17 +++ .../com/nvidia/cuvs/spi/CuVSProvider.java | 19 +++ .../nvidia/cuvs/spi/UnsupportedProvider.java | 5 + .../nvidia/cuvs/internal/CagraIndexImpl.java | 13 ++ .../com/nvidia/cuvs/spi/JDKProvider.java | 5 + .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 29 ++++- .../cuvs/lucene/FilterCuVSProvider.java | 5 + .../TestCagraIndexAtAlignedDimensions.java | 115 ++++++++++++++++++ 8 files changed, 204 insertions(+), 4 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index 51403982dd..96e31431f8 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -336,6 +336,23 @@ static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) thro return CuVSProvider.provider().mergeCagraIndexes(indexes, mergeParams); } + /** + * Reports whether the rows of {@code dataset} already sit at the row stride CAGRA requires, which + * is the row length in bytes rounded up to a 16 byte boundary. + * + *

Use it to pick between the two padded dataset factories: a matrix that is already padded has + * to go through {@link #makePaddedDatasetView(CuVSMatrix)}, because cuVS rejects a request to + * copy it into padded storage it already occupies, and one that is not has to go through + * {@link #makePaddedDataset(CuVSMatrix)}. + * + * @param dataset the matrix to inspect + * @return true when the rows are already padded the way CAGRA requires + */ + static boolean isPaddedDataset(CuVSMatrix dataset) { + Objects.requireNonNull(dataset); + return CuVSProvider.provider().isCagraPaddedDataset(dataset); + } + /** * Builder helps configure and create an instance of {@link CagraIndex}. */ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 05bbc4ab0f..670f44eda6 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -188,6 +188,25 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams merg return mergeCagraIndexes(indexes); } + /** + * Reports whether the rows of {@code dataset} already sit at the row stride CAGRA requires, which + * is the row length in bytes rounded up to a 16 byte boundary. + * + *

This is the question that decides which of the two padded dataset factories a caller has to + * use: {@link CagraIndex#makePaddedDatasetView(CuVSMatrix)} for a device matrix that is already at + * that stride, and {@link CagraIndex#makePaddedDataset(CuVSMatrix)} for one that is not. Asking + * for the wrong one is an error rather than an inefficiency, and the stride of a matrix is not + * visible outside this library, so callers cannot answer it for themselves. + * + * @param dataset the matrix to inspect + * @return true when the rows are already padded the way CAGRA requires + * @throws UnsupportedOperationException if this provider cannot answer + */ + default boolean isCagraPaddedDataset(CuVSMatrix dataset) { + throw new UnsupportedOperationException( + "Padded layout detection is not supported by " + getClass().getName()); + } + /** * Creates a device-backed multi-partition filter handle from the pre-packed combined bitset. * Per-partition bit offsets are recomputed inside cuVS from the index sizes. diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 2a9f9ac738..66b0b8cc47 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -85,6 +85,11 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { throw new UnsupportedOperationException(reasons); } + @Override + public boolean isCagraPaddedDataset(CuVSMatrix dataset) { + throw new UnsupportedOperationException(reasons); + } + @Override public CuVSMatrix.Builder newHostMatrixBuilder( long size, long dimensions, CuVSMatrix.DataType dataType) { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 691c99e93f..9506b130cd 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -263,6 +263,19 @@ private static int elementSizeBytes(CuVSMatrix.DataType dataType) { }; } + /** + * True when the matrix row width matches CAGRA's required padded width for its logical column + * count and element type. The stride lives on the internal matrix type, so this is the only place + * that can answer the question; {@link com.nvidia.cuvs.CagraIndex#isPaddedDataset} routes here + * through the provider. + */ + public static boolean isPaddedDataset(CuVSMatrix dataset) { + if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { + throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); + } + return isCagraPaddedLayout(datasetInternal); + } + /** * True when the matrix row width matches CAGRA's required padded width for its * logical column count and element type. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 4d1dbecf0f..9cc4a5499c 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -317,6 +317,11 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams merge return CagraIndexImpl.merge(indexes, mergeParams); } + @Override + public boolean isCagraPaddedDataset(CuVSMatrix dataset) { + return CagraIndexImpl.isPaddedDataset(dataset); + } + @Override public GPUInfoProvider gpuInfoProvider() { return new GPUInfoProviderImpl(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 97fedf924f..aead683dc4 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -206,6 +206,14 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. + // Log it to make it more obvious that this is what is happening. + info( + infoStream, + COMPONENT, + "CAGRA build failed for field \"" + + fieldInfo.name + + "\", falling back to a brute force index: " + + t); Utils.handleThrowableWithIgnore(t, t.getMessage()); indexType = IndexType.BRUTE_FORCE; } @@ -248,10 +256,23 @@ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwab .withDataset(dataset) .withIndexParams(params) .build(); - var deviceVectors = dataset.toDevice(getCuVSResourcesInstance()); - var indexDataset = index.makePaddedDataset(deviceVectors)) { - index.updateDataset(indexDataset); - index.serialize(os); + var deviceVectors = dataset.toDevice(getCuVSResourcesInstance())) { + /* + * cuVS rejects makePaddedDataset for a device matrix whose rows already sit at the required + * stride, and asks for a view over that storage instead. Copying would be pointless there + * anyway, so pick the factory that matches the layout. + */ + if (CagraIndex.isPaddedDataset(deviceVectors)) { + try (var indexDatasetView = index.makePaddedDatasetView(deviceVectors)) { + index.updateDataset(indexDatasetView); + index.serialize(os); + } + } else { + try (var indexDataset = index.makePaddedDataset(deviceVectors)) { + index.updateDataset(indexDataset); + index.serialize(os); + } + } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index b24d35ffa7..1c74649843 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -85,6 +85,11 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] arg0) throws Throwable { return delegate.mergeCagraIndexes(arg0); } + @Override + public boolean isCagraPaddedDataset(CuVSMatrix arg0) { + return delegate.isCagraPaddedDataset(arg0); + } + @Override public GPUInfoProvider gpuInfoProvider() { return delegate.gpuInfoProvider(); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java new file mode 100644 index 0000000000..89d58aa682 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java @@ -0,0 +1,115 @@ +/* + * 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.isSupported; + +import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.junit.Assume; +import org.junit.Test; + +/** + * A CAGRA row is padded to a 16 byte boundary, so a device matrix whose dimension is already a + * multiple of four sits at the required stride. cuVS refuses to build an owning padded copy of such + * a matrix and asks for a view instead, and the writer swallows a failed CAGRA build by falling + * back to a brute force index. The two together are silent: search keeps returning correct results + * while nothing on the GPU is a CAGRA index any more. + * + *

These tests pin the dimensions on both sides of that boundary. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestCagraIndexAtAlignedDimensions extends LuceneTestCase { + + @Test + public void testCagraIsBuiltAtAnAlignedDimension() throws IOException { + // 128 floats is 512 bytes, an exact multiple of the 16 byte CAGRA row alignment. + assertCagraIsBuilt(128); + } + + @Test + public void testCagraIsBuiltAtAnUnalignedDimension() throws IOException { + // 127 floats is not, so the writer has to fall back to an owning padded copy. + assertCagraIsBuilt(127); + } + + /** Indexes a segment of the given dimension and fails if the CAGRA build did not survive it. */ + private void assertCagraIsBuilt(int dimension) throws IOException { + Assume.assumeTrue("Requires a GPU", isSupported()); + + RecordingInfoStream infoStream = new RecordingInfoStream(); + try (Directory directory = newDirectory()) { + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec( + TestUtil.alwaysKnnVectorsFormat( + new CuVS2510GPUVectorsFormat( + new GPUSearchParams.Builder().withIndexType(IndexType.CAGRA).build()))) + .setInfoStream(infoStream); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int i = 0; i < 64; i++) { + float[] vector = new float[dimension]; + for (int d = 0; d < dimension; d++) { + vector[d] = random().nextFloat(); + } + Document doc = new Document(); + doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + } + } + + assertTrue( + "The CAGRA build fell back to brute force at dimension " + + dimension + + ", messages: " + + infoStream.messages(), + infoStream.cagraBuildFailures().isEmpty()); + } + + /** An InfoStream that keeps the messages, so that a test can tell which index type was built. */ + private static class RecordingInfoStream extends InfoStream { + + private final List messages = Collections.synchronizedList(new ArrayList<>()); + + @Override + public void message(String component, String message) { + messages.add(component + ": " + message); + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + + List messages() { + synchronized (messages) { + return List.copyOf(messages); + } + } + + List cagraBuildFailures() { + return messages().stream().filter(message -> message.contains("CAGRA build failed")).toList(); + } + } +}