Add GPUPlugin for building vector indices on GPU

.buildkite/pipelines/cuvs-snapshot/run-tests.yml

@@ -0,0 +1,21 @@
+steps:
+  - label: "{{matrix.GRADLE_TASK}}"
+    command: .buildkite/scripts/cuvs-snapshot/run-gradle.sh -Dbwc.checkout.align=true -Dorg.elasticsearch.build.cache.push=true -Dignore.tests.seed -Dscan.capture-file-fingerprints {{matrix.GRADLE_TASK}}
+    timeout_in_minutes: 300
+    agents:
+      provider: gcp
+      image: family/elasticsearch-ubuntu-2404-nvidia
+      machineType: g2-standard-32
+      buildDirectory: /dev/shm/bk
+      zones: us-central1-b,us-central1-c
+    env:
+      GRADLE_TASK: "{{matrix.GRADLE_TASK}}"
+    matrix:
+      setup:
+        GRADLE_TASK:
+          - checkPart1
+          - checkPart2
+          - checkPart3
+          - checkPart4
+          - checkPart5
+          - checkPart6

.buildkite/pipelines/cuvs-snapshot/update-snapshot.yml

@@ -0,0 +1,14 @@
+steps:
+  - label: "Smoke test and update new cuVS snapshot"
+    command: .buildkite/scripts/cuvs-snapshot/update-current-snapshot-version.sh
+    agents:
+      provider: gcp
+      image: family/elasticsearch-ubuntu-2404-nvidia
+      machineType: g2-standard-16
+      zones: us-central1-b,us-central1-c
+      diskSizeGb: 150
+  - wait: ~
+  - trigger: "elasticsearch-cuvs-run-tests"
+    build:
+      branch: "${BUILDKITE_BRANCH}"
+    async: true

.buildkite/pipelines/pull-request/.defaults.yml

@@ -1,5 +1,6 @@
 config:
-  skip-labels: ">test-mute"
+  skip-labels:
+    - ">test-mute"
   excluded-regions:
     - ^docs/.*
     - ^x-pack/docs/.*

.buildkite/pipelines/pull-request/bwc-snapshots.yml

@@ -78,4 +78,3 @@ steps:
           image: family/elasticsearch-ubuntu-2004
           machineType: n1-standard-32
           buildDirectory: /dev/shm/bk
-

.buildkite/pipelines/pull-request/gpu.yml

@@ -0,0 +1,28 @@
+config:
+  allow-labels: test-gpu
+  skip-labels:
+    - ">test-mute"
+steps:
+  - group: gpu-tests
+    steps:
+      - label: "{{matrix.GRADLE_TASK}} / gpu-tests"
+        key: "packaging-tests-unix"
+        command: .buildkite/scripts/cuvs-snapshot/run-gradle.sh -Dbwc.checkout.align=true -Dorg.elasticsearch.build.cache.push=true -Dignore.tests.seed -Dscan.capture-file-fingerprints {{matrix.GRADLE_TASK}}
+        timeout_in_minutes: 300
+        agents:
+          provider: gcp
+          image: family/elasticsearch-ubuntu-2404-nvidia
+          machineType: g2-standard-32
+          buildDirectory: /dev/shm/bk
+          zones: us-central1-b,us-central1-c
+        env:
+          GRADLE_TASK: "{{matrix.GRADLE_TASK}}"
+        matrix:
+          setup:
+            GRADLE_TASK:
+              - checkPart1
+              - checkPart2
+              - checkPart3
+              - checkPart4
+              - checkPart5
+              - checkPart6

.buildkite/scripts/cuvs-snapshot/configure.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+
+set -euo pipefail
+
+if [[ -f /etc/profile.d/elastic-nvidia.sh ]]; then
+  export JAVA_HOME="$HOME/.java/openjdk24"
+  export PATH="$JAVA_HOME/bin:$PATH"
+
+  # Setup LD_LIBRARY_PATH, PATH
+
+  export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}"
+  source /etc/profile.d/elastic-nvidia.sh
+fi
+
+# Not running this before the tests results in an error when running the tests
+# No idea why...
+nvidia-smi
+
+CURRENT_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+ELASTICSEARCH_REPO_DIR="$(cd "$CURRENT_SCRIPT_DIR/../../.." && pwd)"
+
+CUVS_SNAPSHOT_VERSION="${CUVS_SNAPSHOT_VERSION:-$(cat "$CURRENT_SCRIPT_DIR"/current-snapshot-version)}"
+CUVS_ARCHIVE="cuvs-$CUVS_SNAPSHOT_VERSION.tar.gz"
+CUVS_URL="https://storage.googleapis.com/elasticsearch-cuvs-snapshots/$CUVS_ARCHIVE"
+
+CUVS_WORKSPACE=${CUVS_WORKSPACE:-$(cd "$(mktemp -d)")}
+CUVS_DIR="$(pwd)/cuvs-$CUVS_SNAPSHOT_VERSION"
+
+curl -O "$CUVS_URL"
+tar -xzf "$CUVS_ARCHIVE"
+
+CUVS_VERSION=$(cd "$CUVS_DIR/cuvs-java/target" && mvn help:evaluate -Dexpression=project.version -q -DforceStdout)
+
+LD_LIBRARY_PATH=$(echo "$LD_LIBRARY_PATH" | tr ':' '\n' | grep -v "libcuvs/linux-x64" | tr '\n' ':' | sed 's/:$//')
+LD_LIBRARY_PATH="$CUVS_DIR/libcuvs/linux-x64:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH
+
+cd "$CUVS_DIR/cuvs-java/target"
+mvn install:install-file -Dfile="cuvs-java-$CUVS_VERSION.jar" -DartifactId=elastic-cuvs-java -DgeneratePom=true
+
+cd "$ELASTICSEARCH_REPO_DIR"
+PLUGIN_GRADLE_FILE=x-pack/plugin/gpu/build.gradle
+sed -i "s|implementation 'com.nvidia.cuvs:elastic-cuvs-java:.*'|implementation 'com.nvidia.cuvs:elastic-cuvs-java:$CUVS_VERSION'|" "$PLUGIN_GRADLE_FILE"

.buildkite/scripts/cuvs-snapshot/current-snapshot-version

@@ -0,0 +1 @@
+fdb8bfb8

.buildkite/scripts/cuvs-snapshot/run-gradle.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+set -euo pipefail
+
+source .buildkite/scripts/cuvs-snapshot/configure.sh
+
+cd "$WORKSPACE"
+
+.ci/scripts/run-gradle.sh "$@"

.buildkite/scripts/cuvs-snapshot/update-current-snapshot-version.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+
+set -euo pipefail
+
+SNAPSHOT_VERSION_FILE=.buildkite/scripts/cuvs-snapshot/current-snapshot-version
+BRANCH_TO_UPDATE="${BRANCH_TO_UPDATE:-${BUILDKITE_BRANCH:-cuvs-snapshot}}"
+
+if [[ -z "${CUVS_SNAPSHOT_VERSION:-}" ]]; then
+  echo "CUVS_SNAPSHOT_VERSION not set. Set this to update the current snapshot version."
+  exit 1
+fi
+
+if [[ "$CUVS_SNAPSHOT_VERSION" == "$(cat $SNAPSHOT_VERSION_FILE)" ]]; then
+  echo "Current snapshot version already set to '$CUVS_SNAPSHOT_VERSION'. No need to update."
+  exit 0
+fi
+
+echo "--- Configuring libcuvs/cuvs-java"
+source .buildkite/scripts/cuvs-snapshot/configure.sh
+
+if [[ "${SKIP_TESTING:-}" != "true" ]]; then
+  echo "--- Testing snapshot before updating"
+  ./gradlew -Druntime.java=24 :x-pack:plugin:gpu:yamlRestTest -S
+fi
+
+echo "--- Updating snapshot"
+
+echo "$CUVS_SNAPSHOT_VERSION" > "$SNAPSHOT_VERSION_FILE"
+
+CURRENT_SHA="$(gh api "/repos/elastic/elasticsearch/contents/$SNAPSHOT_VERSION_FILE?ref=$BRANCH_TO_UPDATE" | jq -r .sha)" || true
+
+gh api -X PUT "/repos/elastic/elasticsearch/contents/$SNAPSHOT_VERSION_FILE" \
+  -f branch="$BRANCH_TO_UPDATE" \
+  -f message="Update cuvs snapshot version to $CUVS_VERSION" \
+  -f content="$(base64 -w 0 "$WORKSPACE/$SNAPSHOT_VERSION_FILE")" \
+  -f sha="$CURRENT_SHA"

docs/changelog/135545.yaml

@@ -0,0 +1,5 @@
+pr: 135545
+summary: Add GPUPlugin for indexing vectors on GPU
+area: Vector Search
+type: feature
+issues: []

gradle/verification-metadata.xml

@@ -1199,6 +1199,11 @@
             <sha256 value="64fab42f17bf8e0efb193dd34da716ef7abb7515234036119df1776b808dc066" origin="Generated by Gradle"/>
          </artifact>
       </component>
+      <component group="com.nvidia.cuvs" name="cuvs-java" version="25.10.0">
+         <artifact name="cuvs-java-25.10.0.jar">
+            <sha256 value="2d4cca3b6b6c7c4d3c1a2f57c00cb55f7a45699536ad356699f32bcea7714539" origin="Generated by Gradle"/>
+         </artifact>
+      </component>
       <component group="com.perforce" name="p4java" version="2015.2.1365273">
          <artifact name="p4java-2015.2.1365273.jar">
             <sha256 value="fbcf286c5863a658b400ec2595ee13c5d49c656f735923088509e3f976ea421e" origin="Generated by Gradle"/>

qa/vector/build.gradle

@@ -22,6 +22,12 @@ tasks.named("dependencyLicenses").configure {
 tasks.named('forbiddenApisMain').configure {
   enabled = false
 }
+repositories {
+  mavenLocal()
+  maven {
+    url = uri("https://storage.googleapis.com/elasticsearch-cuvs-snapshots")
+  }
+}
 
 dependencies {
   api "org.apache.lucene:lucene-core:${versions.lucene}"
@@ -31,6 +37,7 @@ dependencies {
   implementation project(':libs:native')
   implementation project(':libs:logging')
   implementation project(':server')
+  implementation project(':x-pack:plugin:gpu')
 }
 /**
  * Task to run the KnnIndexTester with the provided parameters.

qa/vector/src/main/java/module-info.java

@@ -18,4 +18,5 @@
     requires org.elasticsearch.logging;
     requires java.management;
     requires jdk.management;
+    requires org.elasticsearch.gpu;
 }

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java

@@ -39,6 +39,8 @@
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.gpu.codec.ES92GpuHnswSQVectorsFormat;
+import org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormat;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -76,7 +78,8 @@ public class KnnIndexTester {
     enum IndexType {
         HNSW,
         FLAT,
-        IVF
+        IVF,
+        GPU_HNSW
     }
 
     enum MergePolicyType {
@@ -90,6 +93,8 @@ private static String formatIndexPath(CmdLineArgs args) {
         List<String> suffix = new ArrayList<>();
         if (args.indexType() == IndexType.FLAT) {
             suffix.add("flat");
+        } else if (args.indexType() == IndexType.GPU_HNSW) {
+            suffix.add("gpu_hnsw");
         } else if (args.indexType() == IndexType.IVF) {
             suffix.add("ivf");
             suffix.add(Integer.toString(args.ivfClusterSize()));
@@ -107,6 +112,16 @@ static Codec createCodec(CmdLineArgs args) {
         final KnnVectorsFormat format;
         if (args.indexType() == IndexType.IVF) {
             format = new ES920DiskBBQVectorsFormat(args.ivfClusterSize(), ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER);
+        } else if (args.indexType() == IndexType.GPU_HNSW) {
+            if (args.quantizeBits() == 32) {
+                format = new ES92GpuHnswVectorsFormat();
+            } else if (args.quantizeBits() == 7) {
+                format = new ES92GpuHnswSQVectorsFormat();
+            } else {
+                throw new IllegalArgumentException(
+                    "GPU HNSW index type only supports 7 or 32 bits quantization, but got: " + args.quantizeBits()
+                );
+            }
         } else {
             if (args.quantizeBits() == 1) {
                 if (args.indexType() == IndexType.FLAT) {

server/src/main/java/module-info.java

@@ -400,7 +400,8 @@
             org.elasticsearch.settings.secure,
             org.elasticsearch.serverless.constants,
             org.elasticsearch.serverless.apifiltering,
-            org.elasticsearch.internal.security;
+            org.elasticsearch.internal.security,
+            org.elasticsearch.gpu;
 
     exports org.elasticsearch.telemetry.tracing;
     exports org.elasticsearch.telemetry;
@@ -486,7 +487,7 @@
     exports org.elasticsearch.plugins.internal.rewriter to org.elasticsearch.inference;
     exports org.elasticsearch.lucene.util.automaton;
     exports org.elasticsearch.index.codec.perfield;
-    exports org.elasticsearch.index.codec.vectors to org.elasticsearch.test.knn;
+    exports org.elasticsearch.index.codec.vectors to org.elasticsearch.test.knn, org.elasticsearch.gpu;
     exports org.elasticsearch.index.codec.vectors.es818 to org.elasticsearch.test.knn;
     exports org.elasticsearch.inference.telemetry;
     exports org.elasticsearch.index.codec.vectors.diskbbq to org.elasticsearch.test.knn;

server/src/main/java/org/elasticsearch/index/codec/PerFieldFormatSupplier.java

@@ -120,7 +120,7 @@ public KnnVectorsFormat getKnnVectorsFormatForField(String field) {
         if (mapperService != null) {
             Mapper mapper = mapperService.mappingLookup().getMapper(field);
             if (mapper instanceof DenseVectorFieldMapper vectorMapper) {
-                return vectorMapper.getKnnVectorsFormatForField(knnVectorsFormat);
+                return vectorMapper.getKnnVectorsFormatForField(knnVectorsFormat, mapperService.getIndexSettings());
             }
         }
         return knnVectorsFormat;

server/src/main/java/org/elasticsearch/index/codec/vectors/ES814ScalarQuantizedVectorsFormat.java

@@ -130,7 +130,7 @@ public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException
         );
     }
 
-    static final class ES814ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
+    public static final class ES814ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
 
         final Lucene99ScalarQuantizedVectorsWriter delegate;
 

server/src/main/java/org/elasticsearch/index/mapper/DocumentParser.java

@@ -800,7 +800,8 @@ private static void postProcessDynamicArrayMapping(DocumentParserContext context
             DenseVectorFieldMapper.Builder builder = new DenseVectorFieldMapper.Builder(
                 fieldName,
                 context.indexSettings().getIndexVersionCreated(),
-                IndexSettings.INDEX_MAPPING_EXCLUDE_SOURCE_VECTORS_SETTING.get(context.indexSettings().getSettings())
+                IndexSettings.INDEX_MAPPING_EXCLUDE_SOURCE_VECTORS_SETTING.get(context.indexSettings().getSettings()),
+                context.getVectorFormatProviders()
             );
             builder.dimensions(mappers.size());
             DenseVectorFieldMapper denseVectorFieldMapper = builder.build(builderContext);

server/src/main/java/org/elasticsearch/index/mapper/DocumentParserContext.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.analysis.IndexAnalyzers;
 import org.elasticsearch.index.mapper.MapperService.MergeReason;
+import org.elasticsearch.index.mapper.vectors.VectorsFormatProvider;
 import org.elasticsearch.xcontent.FilterXContentParserWrapper;
 import org.elasticsearch.xcontent.FlatteningXContentParser;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -305,6 +306,10 @@ public final MetadataFieldMapper getMetadataMapper(String mapperName) {
         return mappingLookup.getMapping().getMetadataMapperByName(mapperName);
     }
 
+    public final List<VectorsFormatProvider> getVectorFormatProviders() {
+        return mappingParserContext.getVectorsFormatProviders();
+    }
+
     public final MappingParserContext dynamicTemplateParserContext(DateFormatter dateFormatter) {
         return mappingParserContext.createDynamicTemplateContext(dateFormatter);
     }

server/src/main/java/org/elasticsearch/index/mapper/MapperRegistry.java

@@ -11,11 +11,13 @@
 
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.IndexVersions;
+import org.elasticsearch.index.mapper.vectors.VectorsFormatProvider;
 import org.elasticsearch.plugins.FieldPredicate;
 import org.elasticsearch.plugins.MapperPlugin;
 
 import java.util.Collections;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
 
@@ -32,21 +34,24 @@ public final class MapperRegistry {
     private final Map<String, MetadataFieldMapper.TypeParser> metadataMapperParsers5x;
     private final Function<String, FieldPredicate> fieldFilter;
     private final RootObjectMapperNamespaceValidator namespaceValidator;
+    private final List<VectorsFormatProvider> vectorsFormatProviders;
 
     public MapperRegistry(
         Map<String, Mapper.TypeParser> mapperParsers,
         Map<String, RuntimeField.Parser> runtimeFieldParsers,
         Map<String, MetadataFieldMapper.TypeParser> metadataMapperParsers,
-        Function<String, FieldPredicate> fieldFilter
+        Function<String, FieldPredicate> fieldFilter,
+        List<VectorsFormatProvider> vectorsFormatProviders
     ) {
-        this(mapperParsers, runtimeFieldParsers, metadataMapperParsers, fieldFilter, null);
+        this(mapperParsers, runtimeFieldParsers, metadataMapperParsers, fieldFilter, vectorsFormatProviders, null);
     }
 
     public MapperRegistry(
         Map<String, Mapper.TypeParser> mapperParsers,
         Map<String, RuntimeField.Parser> runtimeFieldParsers,
         Map<String, MetadataFieldMapper.TypeParser> metadataMapperParsers,
         Function<String, FieldPredicate> fieldFilter,
+        List<VectorsFormatProvider> vectorsFormatProviders,
         RootObjectMapperNamespaceValidator namespaceValidator
     ) {
         this.mapperParsers = Collections.unmodifiableMap(new LinkedHashMap<>(mapperParsers));
@@ -62,6 +67,7 @@ public MapperRegistry(
         this.metadataMapperParsers5x = metadata5x;
         this.fieldFilter = fieldFilter;
         this.namespaceValidator = namespaceValidator;
+        this.vectorsFormatProviders = vectorsFormatProviders;
     }
 
     /**
@@ -88,6 +94,10 @@ public RootObjectMapperNamespaceValidator getNamespaceValidator() {
         return namespaceValidator;
     }
 
+    public List<VectorsFormatProvider> getVectorsFormatProviders() {
+        return vectorsFormatProviders;
+    }
+
     /**
      * Return a map of the meta mappers that have been registered. The
      * returned map uses the name of the field as a key.

server/src/main/java/org/elasticsearch/index/mapper/MapperService.java

@@ -246,6 +246,7 @@ public MapperService(
             indexSettings,
             idFieldMapper,
             bitSetProducer,
+            mapperRegistry.getVectorsFormatProviders(),
             mapperRegistry.getNamespaceValidator()
         );
         this.documentParser = new DocumentParser(parserConfiguration, this.mappingParserContextSupplier.get());