gpu hnsw

add gpu indexing flag

add gpu vector storage time measurement

add chmod flag

build hnsw using gpu

scoring speedup

move gl_GlobalInvocationID to mains

glslc with o flag

fix tests

use gpu atomics instead of multiple runs

enable atomics

add timers

profile search

let shader work while points collection

use generations for visited flags

add gpu init timer

avoid entries download

remove gpu call timers for a while

rename to combined graph builder

move gpu graph builder

restore levels in link access

restore links_map

add gpu start condition

use graph builder links memory

remove mut's

list point ids before processing

download-upload links using methods only

cpu graph builder as separate struct

add cpu threads count

parallel cpu build

single threaded cpu preprocess

try to fix glove accuracy

unsafe Send and Sync for gpu builder

copy graph layers builder instead of moving

remove obsolete mut

fix build

move builders to arc

fix bfb deadlock

move uneccessary clear

run cpu and gpu in separate threads

many vectors storage buffers

debug gpu vector storage

fix vector storage upload bug

add cpu ful graph build

don't use non-ready links for point

try to resolve conflicts while upsert

add indexing unit test for measurements

add test

utilize cpu more efficient

dont clear bhaep

debug gpu runs

try to use vec4 instead of mat4

dump scores count

estimate gpu usage

print usage per run

start using working groups

build shaders for vulkan v1.3

start shaders subgroup

vk instance version 1.3

subgroup vector storage

fix test

nearest heap

nearest heap test

are you happy fmt

add assert to nearest heap test

are you happy fmt

gpu visited flags

start candidates heap

candidates heap push

debug nearest heap

remove subgroupExclusiveMin

fix build

are you happy clippy

remove links dependency from nearest heap

provide input count to nearest test

are you happy fmt

provide ef to shader

are you happy fmt

get subgroup size

move compiled shaders

are you happy fmt

fix workgroups count for vector storage compute test

return true subgroup size

do nearest heap on gpu with partial sorting

remove tmp shader

unite scores and indices

are you happy fmt

gpu nearest sorting

are you happy fmt

fix nearest heap test

links test with subgroups

gpu candidates start

fix candidates sorting

are you happy fmt

fix candidates test on m1 chip

move obsolete and test shaders

start new gpu search context

search context test

fix links uploading

debug hnsw searh on level test

fix gpu hnsw search on level test

move test initialization to separate function

gpu greedy search shader test

test heuristic shader

gpu greedy search test

gpu test heuristic

move greedy search into searcher

start insertion shader

find hnsw patch on gpu

fix gpu patches test

start gpu graph builder

equivalency test

fix layers check

add timers

are you happy fmt

quality test and fix multithreaded bug

debug tests

are you happy fmt

fix levels count and clear links

fix equivalency test

remove obsolete solution

trivial gpu indexing integration

gpu config and skip conflicts for large hnsw

glsl generic vector storage element type

gpu generic vector storage

enable vulkan features for f16 and u8

force_half_precision option

add force f16 test

fix u8 and f16 unit tests

apply generic shaders in hnsw construction

remove conflicts check; start cpu utilization

cpu prebuild

upload first point links to gpu

mark all points in cpu graph as ready

multithreaded cpu and gpu

fix equivalency test

remove cpu links copy

parallel patches clearing

fix multithreaded env

external layers source

more logs to debug cpu+gpu syncronization

clear gpu processed points count

remove some temp logs

fix cpu stop condition bug

new insert vector shader

greedy search returns only ids

download whole layer links

apply new links on gpu directly

update win shader script

parallel links loading

provide memory amount as setting and estimate groups count

start bq

gpu bq test

fix tests

shaders invocations refactor

revert max groups settings

gix build after rebase

fix bug after rebase

upload links to gpu in cpu thread

fix greedy search iteration

atomic checker for conflicts

dump graph changes for debugging

move iterator to separate file

start combined builder refactor

start cpu builder as separate struct

separate gpu and cpu construction

fix tests

fix atomics count bug

upload links to gpu as separate fn

glue cpu and gpu

gpu thread

link points on cpu while gpu is busy

fix gpu bq scoring test

use xor instead of calc it

reallocate with less groups if not enough gpu memory

update compiled shaders

don't use llvm gpu emulator

debug panic if no gpu

remove tmp debug dumping

dockerfile for gpu

customize candidates count

fix gpu storage size

runtime shader compilation

dim as define

nearest heap params as macro

candidates and links capacities as macro

parallel bitonic sort

use bubble sort too

nearest heap as shared

limited shared candidates heap

combine buble and bitonic sorts

one-argumented similarity

cache vector for scoring

dont check candidates overflow

do greedy before search

visited hashtable

Revert "visited hashtable"

This reverts commit 7540d04d3859d02f17908c874927d18076f3eb8e.

use less flushes

bulk sum calculation

Revert "bulk sum calculation"

This reverts commit 6ef78ceb8d6ba7d9bf2077804548df162052c749.

load vector in one read

Revert "load vector in one read"

This reverts commit ea980d0399b90f44638590d40fdbfb1d6ffa9504.

fix build after rebase

fix build after rebase

sq support

simplify bq uploading using constructed quantization

metrics

alignment depends on subgroup size

are you happy fmt

exact flag

use lazy static for device access

remove prints

test reallocation factor

visited flags capacity as a compile time constant

build gpu hnsw for payload blocks

add dynamically sized subgroup size support

fix build

update rust version for nvidia docker

dynamically sized subgroups

start device filter

test device filter

are you happy clippy

devices manager

provide device instead of flag, better locking

partial cpu permit release

free cpus when gpu enabled

use cpu is waiting is disabled

create device only when gpu is enabled

cpu permit release all cases

parallel indices count

provide queue index

small refactor of shader builder

are you happy fmt

quantization params as separate struct

pq refactor

fmt

remove upload_mapped_ptr from gpu buffer

refactor gpu buffer

check buffer range

buffer refactor

refactor pipeline

move shader compiler to instance

refactor context

more refactor

no unwraps in gpu crate

clippy

apply pr changes

are you happy fmt

dont use shaderc in segment

more gpu tests

pq bq tests

provide bq option

fix bq for small dim

bind pq buffers

pq more complex test

cover all storage types by tests

test multi and vectors as iterator

multivectors upload data

multivectors shader

fix multivector tests

refactor shader builder

move storage tests into separate file

move quantization to separate file

refactor gpu vector storage creation

refactor gpu quantization api

dont convert to float when unnecessary

more comments

small upload vectors refactor

fix build after rebase

fix tests after rebase

different vk queue priorities capacity

vectors upload stopper

amd dockerfile

fix build after rebase

Cargo.lock

@@ -4945,6 +4945,7 @@ dependencies = [
  "prost 0.11.9",
  "pyroscope",
  "pyroscope_pprofrs",
+ "quantization",
  "raft",
  "raft-proto",
  "rand 0.8.5",
@@ -5973,6 +5974,7 @@ dependencies = [
  "is_sorted",
  "issues",
  "itertools 0.13.0",
+ "lazy_static",
  "log",
  "macro_rules_attribute",
  "memmap2 0.9.5",

Cargo.toml

@@ -17,7 +17,7 @@ default-run = "qdrant"
 workspace = true
 
 [features]
-default = ["web", "parking_lot"]
+default = ["web", "parking_lot", "gpu"]
 web = ["actix-web"]
 multiling-chinese = ["segment/multiling-chinese"]
 multiling-japanese = ["segment/multiling-japanese"]
@@ -100,6 +100,7 @@ common = { path = "lib/common/common" }
 cancel = { path = "lib/common/cancel" }
 memory = { path = "lib/common/memory" }
 issues = { path = "lib/common/issues" }
+quantization = { path = "lib/quantization" }
 segment = { path = "lib/segment" }
 collection = { path = "lib/collection" }
 storage = { path = "lib/storage" }

DockerfileGpuAmd

@@ -0,0 +1,192 @@
+# Cross-compiling using Docker multi-platform builds/images and `xx`.
+#
+# https://docs.docker.com/build/building/multi-platform/
+# https://github.com/tonistiigi/xx
+FROM --platform=${BUILDPLATFORM:-linux/amd64} tonistiigi/xx AS xx
+
+# Utilizing Docker layer caching with `cargo-chef`.
+#
+# https://www.lpalmieri.com/posts/fast-rust-docker-builds/
+FROM --platform=${BUILDPLATFORM:-linux/amd64} lukemathwalker/cargo-chef:latest-rust-1.82.0 AS chef
+
+
+FROM chef AS planner
+WORKDIR /qdrant
+COPY . .
+RUN cargo chef prepare --recipe-path recipe.json
+
+
+FROM chef AS builder
+WORKDIR /qdrant
+
+COPY --from=xx / /
+
+# Relative order of `ARG` and `RUN` commands in the Dockerfile matters.
+#
+# If you pass a different `ARG` to `docker build`, it would invalidate Docker layer cache
+# for the next steps. (E.g., the following steps may depend on a new `ARG` value, so Docker would
+# have to re-execute them instead of using a cached layer from a previous run.)
+#
+# Steps in this stage are ordered in a way that should maximize Docker layer cache utilization,
+# so, please, don't reorder them without prior consideration. 🥲
+
+RUN apt-get update \
+    && apt-get install -y clang lld cmake protobuf-compiler jq \
+    && rustup component add rustfmt
+
+# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility.
+#
+# https://github.com/docker/buildx/issues/510
+ARG BUILDPLATFORM
+ENV BUILDPLATFORM=${BUILDPLATFORM:-linux/amd64}
+
+ARG MOLD_VERSION=2.31.0
+
+RUN case "$BUILDPLATFORM" in \
+        */amd64 ) PLATFORM=x86_64 ;; \
+        */arm64 | */arm64/* ) PLATFORM=aarch64 ;; \
+        * ) echo "Unexpected BUILDPLATFORM '$BUILDPLATFORM'" >&2; exit 1 ;; \
+    esac; \
+    \
+    mkdir -p /opt/mold; \
+    cd /opt/mold; \
+    \
+    TARBALL="mold-$MOLD_VERSION-$PLATFORM-linux.tar.gz"; \
+    curl -sSLO "https://github.com/rui314/mold/releases/download/v$MOLD_VERSION/$TARBALL"; \
+    tar -xf "$TARBALL" --strip-components 1; \
+    rm "$TARBALL"
+
+# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility.
+#
+# https://github.com/docker/buildx/issues/510
+ARG TARGETPLATFORM
+ENV TARGETPLATFORM=${TARGETPLATFORM:-linux/amd64}
+
+RUN xx-apt-get install -y pkg-config gcc g++ libc6-dev libunwind-dev
+
+# Select Cargo profile (e.g., `release`, `dev` or `ci`)
+ARG PROFILE=release
+
+# Enable crate features
+ARG FEATURES
+
+# Pass custom `RUSTFLAGS` (e.g., `--cfg tokio_unstable` to enable Tokio tracing/`tokio-console`)
+ARG RUSTFLAGS
+
+# Select linker (e.g., `mold`, `lld` or an empty string for the default linker)
+ARG LINKER=mold
+
+COPY --from=planner /qdrant/recipe.json recipe.json
+# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`!
+#
+# https://github.com/tonistiigi/xx/issues/107
+# https://github.com/tonistiigi/xx/pull/108
+RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \
+    PATH="$PATH:/opt/mold/bin" \
+    RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \
+    xx-cargo chef cook --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --recipe-path recipe.json
+
+COPY . .
+# Include git commit into Qdrant binary during build
+ARG GIT_COMMIT_ID
+# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`!
+#
+# https://github.com/tonistiigi/xx/issues/107
+# https://github.com/tonistiigi/xx/pull/108
+RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \
+    PATH="$PATH:/opt/mold/bin" \
+    RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \
+    xx-cargo build --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --bin qdrant \
+    && PROFILE_DIR=$(if [ "$PROFILE" = dev ]; then echo debug; else echo $PROFILE; fi) \
+    && mv target/$(xx-cargo --print-target-triple)/$PROFILE_DIR/qdrant /qdrant/qdrant
+
+# Download and extract web UI
+RUN mkdir /static && STATIC_DIR=/static ./tools/sync-web-ui.sh
+
+
+FROM rocm/dev-ubuntu-22.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libglvnd0 \
+    libgl1 \
+    libglx0 \
+    libegl1  \
+    libgles2  \
+    libxcb1-dev \
+    wget \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN apt-get update \
+    && apt-get install -y \
+    libxext6 \
+    libvulkan1 \
+    libvulkan-dev \
+    vulkan-tools
+
+# Install additional packages into the container.
+# E.g., the debugger of choice: gdb/gdbserver/lldb.
+ARG PACKAGES
+
+RUN apt-get install -y --no-install-recommends ca-certificates tzdata libunwind8 $PACKAGES \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy Qdrant source files into the container. Useful for debugging.
+#
+# To enable, set `SOURCES` to *any* non-empty string. E.g., 1/true/enable/whatever.
+# (Note, that *any* non-empty string would work, so 0/false/disable would enable the option as well.)
+ARG SOURCES
+
+# Dockerfile does not support conditional `COPY` instructions (e.g., it's impossible to do something
+# like `if [ -n "$SOURCES" ]; then COPY ...; fi`), so we *hack* conditional `COPY` by abusing
+# parameter expansion and `COPY` wildcards support. 😎
+
+ENV DIR=${SOURCES:+/qdrant/src}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=${SOURCES:+/qdrant/lib}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=${SOURCES:+/usr/local/cargo/registry/src}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=${SOURCES:+/usr/local/cargo/git/checkouts}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=
+
+ARG APP=/qdrant
+
+COPY --from=builder /qdrant/qdrant "$APP"/qdrant
+COPY --from=builder /qdrant/config "$APP"/config
+COPY --from=builder /qdrant/tools/entrypoint.sh "$APP"/entrypoint.sh
+COPY --from=builder /static "$APP"/static
+
+WORKDIR "$APP"
+
+ARG USER_ID=0
+
+RUN if [ "$USER_ID" != 0 ]; then \
+        groupadd --gid "$USER_ID" qdrant; \
+        useradd --uid "$USER_ID" --gid "$USER_ID" -m qdrant; \
+        mkdir -p "$APP"/storage "$APP"/snapshots; \
+        chown -R "$USER_ID:$USER_ID" "$APP"; \
+    fi
+
+USER "$USER_ID:$USER_ID"
+
+ENV TZ=Etc/UTC \
+    RUN_MODE=production
+
+EXPOSE 6333
+EXPOSE 6334
+
+LABEL org.opencontainers.image.title="Qdrant"
+LABEL org.opencontainers.image.description="Official Qdrant image"
+LABEL org.opencontainers.image.url="https://qdrant.com/"
+LABEL org.opencontainers.image.documentation="https://qdrant.com/docs"
+LABEL org.opencontainers.image.source="https://github.com/qdrant/qdrant"
+LABEL org.opencontainers.image.vendor="Qdrant"
+
+CMD ["./entrypoint.sh"]

DockerfileGpuNvidia

@@ -0,0 +1,195 @@
+# Cross-compiling using Docker multi-platform builds/images and `xx`.
+#
+# https://docs.docker.com/build/building/multi-platform/
+# https://github.com/tonistiigi/xx
+FROM --platform=${BUILDPLATFORM:-linux/amd64} tonistiigi/xx AS xx
+
+# Utilizing Docker layer caching with `cargo-chef`.
+#
+# https://www.lpalmieri.com/posts/fast-rust-docker-builds/
+FROM --platform=${BUILDPLATFORM:-linux/amd64} lukemathwalker/cargo-chef:latest-rust-1.82.0 AS chef
+
+
+FROM chef AS planner
+WORKDIR /qdrant
+COPY . .
+RUN cargo chef prepare --recipe-path recipe.json
+
+
+FROM chef AS builder
+WORKDIR /qdrant
+
+COPY --from=xx / /
+
+# Relative order of `ARG` and `RUN` commands in the Dockerfile matters.
+#
+# If you pass a different `ARG` to `docker build`, it would invalidate Docker layer cache
+# for the next steps. (E.g., the following steps may depend on a new `ARG` value, so Docker would
+# have to re-execute them instead of using a cached layer from a previous run.)
+#
+# Steps in this stage are ordered in a way that should maximize Docker layer cache utilization,
+# so, please, don't reorder them without prior consideration. 🥲
+
+RUN apt-get update \
+    && apt-get install -y clang lld cmake protobuf-compiler jq \
+    && rustup component add rustfmt
+
+# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility.
+#
+# https://github.com/docker/buildx/issues/510
+ARG BUILDPLATFORM
+ENV BUILDPLATFORM=${BUILDPLATFORM:-linux/amd64}
+
+ARG MOLD_VERSION=2.31.0
+
+RUN case "$BUILDPLATFORM" in \
+        */amd64 ) PLATFORM=x86_64 ;; \
+        */arm64 | */arm64/* ) PLATFORM=aarch64 ;; \
+        * ) echo "Unexpected BUILDPLATFORM '$BUILDPLATFORM'" >&2; exit 1 ;; \
+    esac; \
+    \
+    mkdir -p /opt/mold; \
+    cd /opt/mold; \
+    \
+    TARBALL="mold-$MOLD_VERSION-$PLATFORM-linux.tar.gz"; \
+    curl -sSLO "https://github.com/rui314/mold/releases/download/v$MOLD_VERSION/$TARBALL"; \
+    tar -xf "$TARBALL" --strip-components 1; \
+    rm "$TARBALL"
+
+# `ARG`/`ENV` pair is a workaround for `docker build` backward-compatibility.
+#
+# https://github.com/docker/buildx/issues/510
+ARG TARGETPLATFORM
+ENV TARGETPLATFORM=${TARGETPLATFORM:-linux/amd64}
+
+RUN xx-apt-get install -y pkg-config gcc g++ libc6-dev libunwind-dev
+
+# Select Cargo profile (e.g., `release`, `dev` or `ci`)
+ARG PROFILE=release
+
+# Enable crate features
+ARG FEATURES
+
+# Pass custom `RUSTFLAGS` (e.g., `--cfg tokio_unstable` to enable Tokio tracing/`tokio-console`)
+ARG RUSTFLAGS
+
+# Select linker (e.g., `mold`, `lld` or an empty string for the default linker)
+ARG LINKER=mold
+
+COPY --from=planner /qdrant/recipe.json recipe.json
+# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`!
+#
+# https://github.com/tonistiigi/xx/issues/107
+# https://github.com/tonistiigi/xx/pull/108
+RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \
+    PATH="$PATH:/opt/mold/bin" \
+    RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \
+    xx-cargo chef cook --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --recipe-path recipe.json
+
+COPY . .
+# Include git commit into Qdrant binary during build
+ARG GIT_COMMIT_ID
+# `PKG_CONFIG=...` is a workaround for `xx-cargo` bug for crates using `pkg-config`!
+#
+# https://github.com/tonistiigi/xx/issues/107
+# https://github.com/tonistiigi/xx/pull/108
+RUN PKG_CONFIG="/usr/bin/$(xx-info)-pkg-config" \
+    PATH="$PATH:/opt/mold/bin" \
+    RUSTFLAGS="${LINKER:+-C link-arg=-fuse-ld=}$LINKER $RUSTFLAGS" \
+    xx-cargo build --profile $PROFILE ${FEATURES:+--features} $FEATURES --features=stacktrace --bin qdrant \
+    && PROFILE_DIR=$(if [ "$PROFILE" = dev ]; then echo debug; else echo $PROFILE; fi) \
+    && mv target/$(xx-cargo --print-target-triple)/$PROFILE_DIR/qdrant /qdrant/qdrant
+
+# Download and extract web UI
+RUN mkdir /static && STATIC_DIR=/static ./tools/sync-web-ui.sh
+
+
+FROM nvidia/opengl:1.0-glvnd-devel-ubuntu22.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libglvnd0 \
+    libgl1 \
+    libglx0 \
+    libegl1  \
+    libgles2  \
+    libxcb1-dev \
+    wget \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN apt-get update \
+    && apt-get install -y \
+    libxext6 \
+    libvulkan1 \
+    libvulkan-dev \
+    vulkan-tools
+
+COPY nvidia_icd.json /etc/vulkan/icd.d/
+ENV NVIDIA_DRIVER_CAPABILITIES compute,graphics,utility
+
+# Install additional packages into the container.
+# E.g., the debugger of choice: gdb/gdbserver/lldb.
+ARG PACKAGES
+
+RUN apt-get install -y --no-install-recommends ca-certificates tzdata libunwind8 $PACKAGES \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy Qdrant source files into the container. Useful for debugging.
+#
+# To enable, set `SOURCES` to *any* non-empty string. E.g., 1/true/enable/whatever.
+# (Note, that *any* non-empty string would work, so 0/false/disable would enable the option as well.)
+ARG SOURCES
+
+# Dockerfile does not support conditional `COPY` instructions (e.g., it's impossible to do something
+# like `if [ -n "$SOURCES" ]; then COPY ...; fi`), so we *hack* conditional `COPY` by abusing
+# parameter expansion and `COPY` wildcards support. 😎
+
+ENV DIR=${SOURCES:+/qdrant/src}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=${SOURCES:+/qdrant/lib}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=${SOURCES:+/usr/local/cargo/registry/src}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=${SOURCES:+/usr/local/cargo/git/checkouts}
+COPY --from=builder ${DIR:-/null?} $DIR/
+
+ENV DIR=
+
+ARG APP=/qdrant
+
+COPY --from=builder /qdrant/qdrant "$APP"/qdrant
+COPY --from=builder /qdrant/config "$APP"/config
+COPY --from=builder /qdrant/tools/entrypoint.sh "$APP"/entrypoint.sh
+COPY --from=builder /static "$APP"/static
+
+WORKDIR "$APP"
+
+ARG USER_ID=0
+
+RUN if [ "$USER_ID" != 0 ]; then \
+        groupadd --gid "$USER_ID" qdrant; \
+        useradd --uid "$USER_ID" --gid "$USER_ID" -m qdrant; \
+        mkdir -p "$APP"/storage "$APP"/snapshots; \
+        chown -R "$USER_ID:$USER_ID" "$APP"; \
+    fi
+
+USER "$USER_ID:$USER_ID"
+
+ENV TZ=Etc/UTC \
+    RUN_MODE=production
+
+EXPOSE 6333
+EXPOSE 6334
+
+LABEL org.opencontainers.image.title="Qdrant"
+LABEL org.opencontainers.image.description="Official Qdrant image"
+LABEL org.opencontainers.image.url="https://qdrant.com/"
+LABEL org.opencontainers.image.documentation="https://qdrant.com/docs"
+LABEL org.opencontainers.image.source="https://github.com/qdrant/qdrant"
+LABEL org.opencontainers.image.vendor="Qdrant"
+
+CMD ["./entrypoint.sh"]

config/config.yaml

@@ -347,3 +347,6 @@ tls:
   # Only works for HTTPS endpoints. Does not support gRPC (and intra-cluster communication).
   # If `null` - TTL is disabled.
   cert_ttl: 3600
+
+gpu:
+  indexing: false

lib/common/common/src/cpu.rs

@@ -194,6 +194,19 @@ impl CpuPermit {
     pub fn release(&mut self) {
         self.permit.take();
     }
+
+    /// Partial release CPU permit, giving them back to the semaphore.
+    pub fn release_count(&mut self, release_count: u32) {
+        if self.num_cpus > release_count {
+            self.num_cpus -= release_count;
+            let permit = self.permit.take();
+            self.permit = permit
+                .map(|mut permit| permit.split(self.num_cpus as usize))
+                .flatten();
+        } else {
+            self.release();
+        }
+    }
 }
 
 impl Drop for CpuPermit {

lib/segment/Cargo.toml

@@ -12,6 +12,7 @@ edition = "2021"
 workspace = true
 
 [features]
+default = ["gpu"]
 multiling-chinese = ["charabia/chinese-segmentation", "charabia/chinese-normalization", "charabia/chinese-normalization-pinyin"]
 multiling-japanese = ["charabia/japanese"]
 multiling-korean = ["charabia/korean"]
@@ -85,6 +86,8 @@ is_sorted = "0.1.1"
 strum = { workspace = true }
 byteorder = { workspace = true }
 zerocopy = { workspace = true }
+lazy_static = "1.5.0"
+env_logger = "0.11"
 
 sysinfo = "0.32"
 charabia = { version = "0.9.1", default-features = false, features = ["greek", "hebrew", "thai"] }

lib/segment/benches/multi_vector_search.rs

@@ -85,6 +85,7 @@ fn multi_vector_search_benchmark(c: &mut Criterion) {
         payload_index: segment.payload_index.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/src/index/hnsw_index/gpu/batched_points.rs

@@ -0,0 +1,159 @@
+use std::ops::Range;
+use std::sync::atomic::AtomicU32;
+
+use ahash::HashSet;
+use common::types::PointOffsetType;
+
+use crate::common::operation_error::OperationResult;
+
+#[derive(Debug)]
+pub struct PointLinkingData {
+    pub point_id: PointOffsetType,
+    pub level: usize,
+    pub batch_index: usize,
+    pub entry: AtomicU32,
+}
+
+#[derive(Debug)]
+pub struct Batch<'a> {
+    pub points: &'a [PointLinkingData],
+    pub level: usize,
+    pub end_index: usize,
+}
+
+pub struct BatchedPoints {
+    pub points: Vec<PointLinkingData>,
+    pub batches: Vec<Range<usize>>,
+    pub ids_by_batches: Vec<HashSet<PointOffsetType>>,
+    pub first_point_id: PointOffsetType,
+    pub levels_count: usize,
+}
+
+impl BatchedPoints {
+    pub fn new(
+        level_fn: impl Fn(PointOffsetType) -> usize,
+        mut ids: Vec<PointOffsetType>,
+        groups_count: usize,
+    ) -> OperationResult<Self> {
+        Self::sort_points_by_level(&level_fn, &mut ids);
+        let first_point_id = ids.remove(0);
+
+        let batches = Self::build_initial_batches(&level_fn, &ids, groups_count);
+
+        let mut points = Vec::with_capacity(ids.len());
+        for (batch_index, batch) in batches.iter().enumerate() {
+            for i in batch.clone() {
+                let point_id = ids[i];
+                let level = level_fn(point_id);
+                points.push(PointLinkingData {
+                    point_id,
+                    level,
+                    batch_index,
+                    entry: first_point_id.into(),
+                });
+            }
+        }
+
+        let ids_by_batches = batches
+            .iter()
+            .map(|batch| {
+                batch
+                    .clone()
+                    .map(|i| ids[i])
+                    .collect::<HashSet<PointOffsetType>>()
+            })
+            .collect();
+
+        Ok(Self {
+            points,
+            batches,
+            ids_by_batches,
+            first_point_id,
+            levels_count: level_fn(first_point_id) + 1,
+        })
+    }
+
+    pub fn is_same_batch(
+        &self,
+        linking_point: &PointLinkingData,
+        other_point: PointOffsetType,
+    ) -> bool {
+        self.ids_by_batches[linking_point.batch_index].contains(&other_point)
+    }
+
+    pub fn iter_batches(&self, skip_count: usize) -> impl Iterator<Item = Batch> {
+        self.batches
+            .iter()
+            .filter(move |batch| batch.end > skip_count)
+            .map(move |batch| {
+                let intersected_batch = std::cmp::max(batch.start, skip_count)..batch.end;
+                let level = self.points[intersected_batch.start].level;
+                Batch {
+                    end_index: intersected_batch.end,
+                    points: &self.points[intersected_batch],
+                    level,
+                }
+            })
+    }
+
+    pub fn sort_points_by_level(
+        level_fn: impl Fn(PointOffsetType) -> usize,
+        ids: &mut [PointOffsetType],
+    ) {
+        ids.sort_by(|&a, &b| {
+            let a_level = level_fn(a);
+            let b_level = level_fn(b);
+            match b_level.cmp(&a_level) {
+                std::cmp::Ordering::Less => std::cmp::Ordering::Less,
+                std::cmp::Ordering::Greater => std::cmp::Ordering::Greater,
+                std::cmp::Ordering::Equal => a.cmp(&b),
+            }
+        });
+    }
+
+    fn build_initial_batches(
+        level_fn: impl Fn(PointOffsetType) -> usize,
+        ids: &[PointOffsetType],
+        groups_count: usize,
+    ) -> Vec<Range<usize>> {
+        let timer = std::time::Instant::now();
+
+        let num_vectors = ids.len();
+        let mut batches: Vec<_> = (0..num_vectors.div_ceil(groups_count))
+            .map(|start| {
+                groups_count * start..std::cmp::min(groups_count * (start + 1), num_vectors)
+            })
+            .collect();
+
+        let mut batch_index = 0usize;
+        while batch_index < batches.len() {
+            let batch = batches[batch_index].clone();
+            let point_id = ids[batch.start];
+            let batch_level = level_fn(point_id);
+            for i in 1..batch.len() {
+                let point_id = ids[batch.start + i];
+                let level = level_fn(point_id);
+                // divide batch by level. all batches must be on the same level
+                if level != batch_level {
+                    let batch1 = batch.start..batch.start + i;
+                    let batch2 = batch.start + i..batch.end;
+                    batches[batch_index] = batch1;
+                    batches.insert(batch_index + 1, batch2);
+                    break;
+                }
+            }
+
+            batch_index += 1;
+        }
+
+        for batch_pair in batches.windows(2) {
+            if batch_pair.len() == 2 {
+                assert_eq!(batch_pair[0].end, batch_pair[1].start);
+            }
+        }
+
+        log::trace!("Initial batchs time: {:?}", timer.elapsed());
+
+        batches
+    }
+}

lib/segment/src/index/hnsw_index/gpu/cpu_level_builder.rs

@@ -0,0 +1,260 @@
+use std::sync::atomic::Ordering;
+use std::sync::Mutex;
+
+use common::types::PointOffsetType;
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use rayon::ThreadPool;
+
+use super::batched_points::{BatchedPoints, PointLinkingData};
+use super::gpu_search_context::GpuRequest;
+use crate::common::operation_error::OperationResult;
+use crate::index::hnsw_index::graph_layers::GraphLayersBase;
+use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder;
+use crate::index::hnsw_index::point_scorer::FilteredScorer;
+use crate::payload_storage::FilterContext;
+use crate::vector_storage::RawScorer;
+
+// Build level on CPU. Returns amount of processed points.
+pub fn build_level_on_cpu<'a>(
+    pool: &ThreadPool,
+    graph_layers_builder: &GraphLayersBuilder,
+    batched_points: &BatchedPoints,
+    level: usize,
+    stop_condition: impl Fn(usize) -> bool + Send + Sync,
+    points_scorer_builder: impl Fn(
+            PointOffsetType,
+        )
+            -> OperationResult<(Box<dyn RawScorer + 'a>, Option<Box<dyn FilterContext + 'a>>)>
+        + Send
+        + Sync,
+) -> OperationResult<usize> {
+    let retry_mutex: Mutex<()> = Default::default();
+    let index = Mutex::new(0usize);
+
+    pool.install(|| {
+        (0..graph_layers_builder.links_layers.len())
+            .into_par_iter()
+            .try_for_each(|_| -> OperationResult<()> {
+                let index = {
+                    let mut locked_index = index.lock().unwrap();
+                    let index = *locked_index;
+                    if index >= batched_points.points.len() {
+                        return Ok(());
+                    }
+
+                    if stop_condition(index) {
+                        return Ok(());
+                    }
+                    *locked_index += 1;
+                    index
+                };
+
+                let linking_point = &batched_points.points[index];
+                if graph_layers_builder.get_point_level(linking_point.point_id) < level {
+                    update_entry_on_cpu(
+                        graph_layers_builder,
+                        linking_point,
+                        level,
+                        &points_scorer_builder,
+                    )
+                } else {
+                    link_point_on_cpu(
+                        &retry_mutex,
+                        graph_layers_builder,
+                        batched_points,
+                        linking_point,
+                        level,
+                        &points_scorer_builder,
+                    )
+                }
+            })
+    })?;
+
+    let processed_count = index.lock().unwrap().min(batched_points.points.len());
+    Ok(processed_count)
+}
+
+fn update_entry_on_cpu<'a>(
+    graph_layers_builder: &GraphLayersBuilder,
+    linking_point: &PointLinkingData,
+    level: usize,
+    points_scorer_builder: impl Fn(
+            PointOffsetType,
+        )
+            -> OperationResult<(Box<dyn RawScorer + 'a>, Option<Box<dyn FilterContext + 'a>>)>
+        + Send
+        + Sync,
+) -> OperationResult<()> {
+    let (raw_scorer, filter_context) = points_scorer_builder(linking_point.point_id)?;
+    let mut points_scorer = FilteredScorer::new(raw_scorer.as_ref(), filter_context.as_deref());
+    let new_entry = graph_layers_builder
+        .search_entry_on_level(
+            linking_point.entry.load(Ordering::Relaxed),
+            level,
+            &mut points_scorer,
+        )
+        .idx;
+    linking_point.entry.store(new_entry, Ordering::Relaxed);
+    raw_scorer.take_hardware_counter().discard_results();
+    Ok(())
+}
+
+fn link_point_on_cpu<'a>(
+    retry_mutex: &Mutex<()>,
+    graph_layers_builder: &GraphLayersBuilder,
+    batched_points: &BatchedPoints,
+    linking_point: &PointLinkingData,
+    level: usize,
+    points_scorer_builder: impl Fn(
+            PointOffsetType,
+        )
+            -> OperationResult<(Box<dyn RawScorer + 'a>, Option<Box<dyn FilterContext + 'a>>)>
+        + Send
+        + Sync,
+) -> OperationResult<()> {
+    let (raw_scorer, filter_context) = points_scorer_builder(linking_point.point_id)?;
+    let mut points_scorer = FilteredScorer::new(raw_scorer.as_ref(), filter_context.as_deref());
+    let (patches, new_entries) = graph_layers_builder.get_patch(
+        GpuRequest {
+            id: linking_point.point_id,
+            entry: linking_point.entry.load(Ordering::Relaxed),
+        },
+        level,
+        &mut points_scorer,
+    );
+
+    let new_entries = if !graph_layers_builder.try_apply_patch(level, patches) {
+        let _retry_guard = retry_mutex.lock().unwrap();
+
+        // try again in single-thread mode
+        let (patches, new_entries) = graph_layers_builder.get_patch(
+            GpuRequest {
+                id: linking_point.point_id,
+                entry: linking_point.entry.load(Ordering::Relaxed),
+            },
+            level,
+            &mut points_scorer,
+        );
+        graph_layers_builder.apply_patch(level, patches);
+        new_entries
+    } else {
+        new_entries
+    };
+
+    raw_scorer.take_hardware_counter().discard_results();
+
+    new_entries
+        .into_iter()
+        .filter(|&new_entry| !batched_points.is_same_batch(linking_point, new_entry))
+        .next()
+        .map(|new_entry| {
+            linking_point.entry.store(new_entry, Ordering::Relaxed);
+        });
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use common::types::PointOffsetType;
+
+    use super::*;
+    use crate::fixtures::index_fixtures::FakeFilterContext;
+    use crate::index::hnsw_index::gpu::batched_points::BatchedPoints;
+    use crate::index::hnsw_index::gpu::create_graph_layers_builder;
+    use crate::index::hnsw_index::gpu::tests::{
+        check_graph_layers_builders_quality, compare_graph_layers_builders,
+        create_gpu_graph_test_data, GpuGraphTestData,
+    };
+    use crate::vector_storage::chunked_vector_storage::VectorOffsetType;
+
+    fn build_cpu_graph(test: &GpuGraphTestData, threads: usize) -> GraphLayersBuilder {
+        let num_vectors = test.graph_layers_builder.links_layers.len();
+        let m = test.graph_layers_builder.m;
+        let m0 = test.graph_layers_builder.m0;
+        let ef = test.graph_layers_builder.ef_construct;
+
+        let batched_points = BatchedPoints::new(
+            |point_id| test.graph_layers_builder.get_point_level(point_id),
+            (0..num_vectors as PointOffsetType).collect(),
+            threads,
+        )
+        .unwrap();
+
+        let graph_layers_builder =
+            create_graph_layers_builder(&batched_points, num_vectors, m, m0, ef, 1).unwrap();
+
+        let pool = rayon::ThreadPoolBuilder::new()
+            .thread_name(|idx| format!("hnsw-build-{idx}"))
+            .num_threads(threads)
+            .build()
+            .unwrap();
+
+        for level in (0..batched_points.levels_count).rev() {
+            build_level_on_cpu(
+                &pool,
+                &graph_layers_builder,
+                &batched_points,
+                level,
+                |_| false,
+                |point_id| {
+                    let fake_filter_context = FakeFilterContext {};
+                    let added_vector = test
+                        .vector_holder
+                        .vectors
+                        .get(point_id as VectorOffsetType)
+                        .to_vec();
+                    let raw_scorer = test
+                        .vector_holder
+                        .get_raw_scorer(added_vector.clone())
+                        .unwrap();
+                    Ok((raw_scorer, Some(Box::new(fake_filter_context))))
+                },
+            )
+            .unwrap();
+        }
+
+        graph_layers_builder
+    }
+
+    #[test]
+    fn test_cpu_hnsw_level_equivalency() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let num_vectors = 1024;
+        let dim = 64;
+        let m = 8;
+        let m0 = 16;
+        let ef = 32;
+
+        let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, 0);
+        let graph_layers_builder = build_cpu_graph(&test, 1);
+
+        compare_graph_layers_builders(&test.graph_layers_builder, &graph_layers_builder);
+    }
+
+    #[test]
+    fn test_cpu_hnsw_level_quality() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let num_vectors = 1024;
+        let dim = 64;
+        let m = 8;
+        let m0 = 16;
+        let ef = 32;
+        let threads = 4;
+        let searches_count = 10;
+        let top = 10;
+
+        let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count);
+        let graph_layers_builder = build_cpu_graph(&test, threads);
+
+        check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.9)
+    }
+}

lib/segment/src/index/hnsw_index/gpu/devices_manager.rs

@@ -0,0 +1,84 @@
+use std::sync::Arc;
+
+use parking_lot::{Mutex, MutexGuard};
+
+use crate::common::operation_error::OperationResult;
+
+pub struct DevicesMaganer {
+    pub devices: Vec<Mutex<Arc<gpu::Device>>>,
+    pub wait_free: bool,
+}
+
+pub struct LockedDevice<'a> {
+    pub locked_device: MutexGuard<'a, Arc<gpu::Device>>,
+}
+
+impl DevicesMaganer {
+    pub fn new(
+        instance: Arc<gpu::Instance>,
+        filter: &str,
+        start_index: usize,
+        count: usize,
+        wait_free: bool,
+        parallel_indexes: usize,
+    ) -> OperationResult<Self> {
+        let filter = filter.to_lowercase();
+        let mut devices = Vec::new();
+        for queue_index in 0..parallel_indexes {
+            devices.extend(
+                instance
+                    .physical_devices()
+                    .iter()
+                    .filter(|device| {
+                        let device_name = device.name.to_lowercase();
+                        device_name.contains(&filter)
+                    })
+                    .skip(start_index)
+                    .take(count)
+                    .filter_map(|physical_device| {
+                        match gpu::Device::new_with_queue_index(
+                            instance.clone(),
+                            physical_device,
+                            queue_index,
+                        ) {
+                            Ok(device) => {
+                                log::info!("Initialized GPU device: {:?}", &physical_device.name);
+                                Some(Mutex::new(device))
+                            }
+                            Err(err) => {
+                                log::error!(
+                                    "Failed to create GPU device: {:?}, error: {:?}",
+                                    &physical_device.name,
+                                    err
+                                );
+                                None
+                            }
+                        }
+                    }),
+            );
+        }
+        Ok(Self { devices, wait_free })
+    }
+
+    pub fn lock_device(&self) -> Option<LockedDevice> {
+        if self.devices.is_empty() {
+            return None;
+        }
+        loop {
+            // TODO(gpu): Add timeout
+            for device in &self.devices {
+                if let Some(guard) = device.try_lock() {
+                    return Some(LockedDevice {
+                        locked_device: guard,
+                    });
+                }
+            }
+
+            if !self.wait_free {
+                return None;
+            }
+
+            std::thread::sleep(std::time::Duration::from_millis(500));
+        }
+    }
+}

lib/segment/src/index/hnsw_index/gpu/gpu_candidates_heap.rs

@@ -0,0 +1,217 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use super::shader_builder::ShaderBuilderParameters;
+
+pub struct GpuCandidatesHeap {
+    pub capacity: usize,
+    pub device: Arc<gpu::Device>,
+}
+
+impl GpuCandidatesHeap {
+    pub fn new(device: Arc<gpu::Device>, capacity: usize) -> gpu::GpuResult<Self> {
+        let ceiled_capacity = capacity.div_ceil(device.subgroup_size()) * device.subgroup_size();
+        Ok(Self {
+            capacity: ceiled_capacity,
+            device,
+        })
+    }
+}
+
+impl ShaderBuilderParameters for GpuCandidatesHeap {
+    fn shader_includes(&self) -> HashMap<String, String> {
+        HashMap::from([(
+            "candidates_heap.comp".to_string(),
+            include_str!("shaders/candidates_heap.comp").to_string(),
+        )])
+    }
+
+    fn shader_defines(&self) -> HashMap<String, Option<String>> {
+        let mut defines = HashMap::new();
+        defines.insert(
+            "CANDIDATES_HEAP_CAPACITY".to_owned(),
+            Some(self.capacity.to_string()),
+        );
+        defines
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BinaryHeap;
+
+    use common::types::{PointOffsetType, ScoredPointOffset};
+    use rand::rngs::StdRng;
+    use rand::{Rng, SeedableRng};
+
+    use super::*;
+    use crate::index::hnsw_index::gpu::shader_builder::ShaderBuilder;
+    use crate::index::hnsw_index::gpu::GPU_TIMEOUT;
+
+    #[repr(C)]
+    struct TestParams {
+        input_counts: u32,
+    }
+
+    #[test]
+    fn test_gpu_candidates_heap() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let capacity = 128;
+        let points_count = 128;
+        let groups_count = 8;
+        let inputs_count = points_count;
+
+        let mut rng = StdRng::seed_from_u64(42);
+        let inputs_data: Vec<ScoredPointOffset> = (0..inputs_count * groups_count)
+            .map(|i| ScoredPointOffset {
+                idx: i as PointOffsetType,
+                score: rng.gen_range(-1.0..1.0),
+            })
+            .collect();
+
+        let debug_messenger = gpu::PanicIfErrorMessenger {};
+        let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap();
+        let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap();
+
+        let gpu_candidates_heap = GpuCandidatesHeap::new(device.clone(), capacity).unwrap();
+
+        let shader = ShaderBuilder::new(device.clone())
+            .with_shader_code(include_str!("shaders/tests/test_candidates_heap.comp"))
+            .with_parameters(&gpu_candidates_heap)
+            .build()
+            .unwrap();
+
+        let input_points_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Input points buffer",
+            gpu::BufferType::Storage,
+            inputs_count * groups_count * std::mem::size_of::<ScoredPointOffset>(),
+        )
+        .unwrap();
+
+        let upload_staging_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Candidates heap upload staging buffer",
+            gpu::BufferType::CpuToGpu,
+            inputs_count * groups_count * std::mem::size_of::<ScoredPointOffset>(),
+        )
+        .unwrap();
+        upload_staging_buffer.upload_slice(&inputs_data, 0).unwrap();
+
+        let mut context = gpu::Context::new(device.clone()).unwrap();
+        context
+            .copy_gpu_buffer(
+                upload_staging_buffer.clone(),
+                input_points_buffer.clone(),
+                0,
+                0,
+                input_points_buffer.size(),
+            )
+            .unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+
+        let test_params_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Test params buffer",
+            gpu::BufferType::Uniform,
+            std::mem::size_of::<TestParams>(),
+        )
+        .unwrap();
+        upload_staging_buffer
+            .upload(
+                &TestParams {
+                    input_counts: inputs_count as u32,
+                },
+                0,
+            )
+            .unwrap();
+        context
+            .copy_gpu_buffer(
+                upload_staging_buffer,
+                test_params_buffer.clone(),
+                0,
+                0,
+                test_params_buffer.size(),
+            )
+            .unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+
+        let scores_output_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Scores output buffer",
+            gpu::BufferType::Storage,
+            inputs_count * groups_count * std::mem::size_of::<ScoredPointOffset>(),
+        )
+        .unwrap();
+
+        let descriptor_set_layout = gpu::DescriptorSetLayout::builder()
+            .add_uniform_buffer(0)
+            .add_storage_buffer(1)
+            .add_storage_buffer(2)
+            .build(device.clone())
+            .unwrap();
+
+        let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone())
+            .add_uniform_buffer(0, test_params_buffer.clone())
+            .add_storage_buffer(1, input_points_buffer.clone())
+            .add_storage_buffer(2, scores_output_buffer.clone())
+            .build()
+            .unwrap();
+
+        let pipeline = gpu::Pipeline::builder()
+            .add_descriptor_set_layout(0, descriptor_set_layout.clone())
+            .add_shader(shader.clone())
+            .build(device.clone())
+            .unwrap();
+
+        context
+            .bind_pipeline(pipeline, &[descriptor_set.clone()])
+            .unwrap();
+        context.dispatch(groups_count, 1, 1).unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+
+        let mut scores_cpu = vec![];
+        for group in 0..groups_count {
+            let mut heap = BinaryHeap::<ScoredPointOffset>::new();
+            for i in 0..inputs_count {
+                let scored_point = inputs_data[group * inputs_count + i];
+                heap.push(scored_point);
+            }
+            while !heap.is_empty() {
+                scores_cpu.push(heap.pop().unwrap());
+            }
+        }
+
+        let download_staging_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Candidates heap download staging buffer",
+            gpu::BufferType::GpuToCpu,
+            scores_output_buffer.size(),
+        )
+        .unwrap();
+        context
+            .copy_gpu_buffer(
+                scores_output_buffer.clone(),
+                download_staging_buffer.clone(),
+                0,
+                0,
+                scores_output_buffer.size(),
+            )
+            .unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+        let mut scores_gpu = vec![ScoredPointOffset::default(); inputs_count * groups_count];
+        download_staging_buffer
+            .download_slice(&mut scores_gpu, 0)
+            .unwrap();
+
+        assert_eq!(scores_gpu, scores_cpu);
+    }
+}

lib/segment/src/index/hnsw_index/gpu/gpu_graph_builder.rs

@@ -0,0 +1,304 @@
+use std::sync::atomic::{AtomicBool, AtomicUsize};
+use std::sync::Arc;
+use std::thread::JoinHandle;
+
+use common::types::PointOffsetType;
+use parking_lot::Mutex;
+use rayon::ThreadPool;
+
+use crate::common::operation_error::OperationResult;
+use crate::index::hnsw_index::gpu::batched_points::BatchedPoints;
+use crate::index::hnsw_index::gpu::cpu_level_builder::build_level_on_cpu;
+use crate::index::hnsw_index::gpu::create_graph_layers_builder;
+use crate::index::hnsw_index::gpu::gpu_level_builder::build_level_on_gpu;
+use crate::index::hnsw_index::gpu::gpu_search_context::GpuSearchContext;
+use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder;
+use crate::payload_storage::FilterContext;
+use crate::vector_storage::quantized::quantized_vectors::QuantizedVectors;
+use crate::vector_storage::{RawScorer, VectorStorageEnum};
+
+#[allow(clippy::too_many_arguments)]
+pub fn build_hnsw_on_gpu<'a>(
+    device: Arc<gpu::Device>,
+    pool: &ThreadPool,
+    reference_graph: &GraphLayersBuilder,
+    max_groups_count: usize,
+    vector_storage: &VectorStorageEnum,
+    quantized_storage: Option<&QuantizedVectors>,
+    entry_points_num: usize,
+    force_half_precision: bool,
+    min_cpu_linked_points_count: usize,
+    exact: bool,
+    ids: Vec<PointOffsetType>,
+    points_scorer_builder: impl Fn(
+            PointOffsetType,
+        )
+            -> OperationResult<(Box<dyn RawScorer + 'a>, Option<Box<dyn FilterContext + 'a>>)>
+        + Send
+        + Sync,
+    stopped: &AtomicBool,
+) -> OperationResult<GraphLayersBuilder> {
+    let num_vectors = reference_graph.links_layers.len();
+    let m = reference_graph.m;
+    let m0 = reference_graph.m0;
+    let ef = reference_graph.ef_construct;
+
+    let gpu_search_context = Arc::new(Mutex::new(GpuSearchContext::new(
+        device,
+        max_groups_count,
+        vector_storage,
+        quantized_storage,
+        m,
+        m0,
+        ef,
+        num_vectors,
+        force_half_precision,
+        exact,
+        stopped,
+    )?));
+
+    let batched_points = Arc::new(BatchedPoints::new(
+        |point_id| reference_graph.get_point_level(point_id),
+        ids,
+        gpu_search_context.lock().groups_count,
+    )?);
+
+    let graph_layers_builder = Arc::new(create_graph_layers_builder(
+        &batched_points,
+        num_vectors,
+        m,
+        m0,
+        ef,
+        entry_points_num,
+    )?);
+
+    let mut gpu_thread_handle: Option<JoinHandle<OperationResult<()>>> = None;
+    let gpu_built_points = Arc::new(AtomicUsize::new(batched_points.points.len()));
+
+    let cpu_stop_condition = |cpu_built_points| {
+        assert_ne!(cpu_built_points, batched_points.points.len());
+
+        // if CPU is faster than GPU, wait for GPU to catch up
+        while cpu_built_points >= gpu_built_points.load(std::sync::atomic::Ordering::Relaxed) {
+            log::trace!(
+                "Waiting for GPU to catch up (CPU ready {}, GPU ready {})",
+                cpu_built_points,
+                gpu_built_points.load(std::sync::atomic::Ordering::Relaxed)
+            );
+            std::thread::sleep(std::time::Duration::from_millis(100));
+        }
+
+        let is_gpu_ready = gpu_built_points.load(std::sync::atomic::Ordering::Relaxed)
+            == batched_points.points.len();
+        let min_required_points_reached = cpu_built_points >= min_cpu_linked_points_count;
+
+        // stop id GPU is ready and min required points are linked
+        min_required_points_reached && is_gpu_ready
+    };
+
+    for level in (0..batched_points.levels_count).rev() {
+        // Start building level on CPU
+        let gpu_start_index = build_level_on_cpu(
+            &pool,
+            &graph_layers_builder,
+            &batched_points,
+            level,
+            &cpu_stop_condition,
+            &points_scorer_builder,
+        )?;
+
+        // CPU is ready, finish prev level GPU thread to continue with this level
+        if let Some(handle) = gpu_thread_handle.take() {
+            handle.join().unwrap()?;
+        }
+
+        // Clear GPU built points counter
+        gpu_built_points.store(0, std::sync::atomic::Ordering::Relaxed);
+
+        log::debug!(
+            "Starting GPU level {}, skipped points {}",
+            level,
+            gpu_start_index
+        );
+
+        let gpu_search_context = gpu_search_context.clone();
+        let graph_layers_builder = graph_layers_builder.clone();
+        let batched_points = batched_points.clone();
+        let gpu_built_points = gpu_built_points.clone();
+        // Continue level building on GPU in separate thread
+        gpu_thread_handle = Some(std::thread::spawn(move || -> OperationResult<()> {
+            let mut gpu_search_context = gpu_search_context.lock();
+            gpu_search_context.upload_links(level, &graph_layers_builder)?;
+            build_level_on_gpu(
+                &mut gpu_search_context,
+                &batched_points,
+                gpu_start_index,
+                level,
+                |count| {
+                    gpu_built_points.store(count, std::sync::atomic::Ordering::Relaxed);
+                },
+            )?;
+            gpu_search_context.download_links(level, &graph_layers_builder)?;
+            Ok(())
+        }));
+    }
+
+    // Wait for the last GPU level
+    if let Some(handle) = gpu_thread_handle.take() {
+        handle.join().unwrap()?;
+    }
+
+    {
+        let gpu_search_context = gpu_search_context.lock();
+        log::debug!(
+            "Gpu graph patches time: {:?}, count {:?}, avg {:?}",
+            &gpu_search_context.patches_timer,
+            gpu_search_context.patches_count,
+            gpu_search_context
+                .patches_timer
+                .checked_div(gpu_search_context.patches_count as u32)
+                .unwrap_or_default(),
+        );
+        log::debug!(
+            "Gpu graph update entries time: {:?}, count {:?}, avg {:?}",
+            &gpu_search_context.updates_timer,
+            gpu_search_context.updates_count,
+            gpu_search_context
+                .updates_timer
+                .checked_div(gpu_search_context.updates_count as u32)
+                .unwrap_or_default(),
+        );
+    }
+
+    Ok(Arc::into_inner(graph_layers_builder).unwrap())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::borrow::Borrow;
+
+    use super::*;
+    use crate::fixtures::index_fixtures::FakeFilterContext;
+    use crate::index::hnsw_index::gpu::tests::{
+        check_graph_layers_builders_quality, compare_graph_layers_builders,
+        create_gpu_graph_test_data, GpuGraphTestData,
+    };
+    use crate::vector_storage::chunked_vector_storage::VectorOffsetType;
+
+    fn build_gpu_graph(
+        test: &GpuGraphTestData,
+        groups_count: usize,
+        min_cpu_linked_points_count: usize,
+        exact: bool,
+    ) -> GraphLayersBuilder {
+        let num_vectors = test.graph_layers_builder.links_layers.len();
+
+        let pool = rayon::ThreadPoolBuilder::new()
+            .thread_name(|idx| format!("hnsw-build-{idx}"))
+            .num_threads(groups_count)
+            .build()
+            .unwrap();
+
+        let debug_messenger = gpu::PanicIfErrorMessenger {};
+        let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap();
+        let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap();
+
+        let ids = (0..num_vectors as PointOffsetType).collect();
+        build_hnsw_on_gpu(
+            device,
+            &pool,
+            &test.graph_layers_builder,
+            groups_count,
+            &test.vector_storage.borrow(),
+            None,
+            1,
+            false,
+            min_cpu_linked_points_count,
+            exact,
+            ids,
+            |point_id| {
+                let fake_filter_context = FakeFilterContext {};
+                let added_vector = test
+                    .vector_holder
+                    .vectors
+                    .get(point_id as VectorOffsetType)
+                    .to_vec();
+                let raw_scorer = test
+                    .vector_holder
+                    .get_raw_scorer(added_vector.clone())
+                    .unwrap();
+                Ok((raw_scorer, Some(Box::new(fake_filter_context))))
+            },
+            &false.into(),
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn test_gpu_hnsw_equivalency() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let num_vectors = 1024;
+        let dim = 64;
+        let m = 8;
+        let m0 = 16;
+        let ef = 32;
+        let min_cpu_linked_points_count = 64;
+
+        let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, 0);
+        let graph_layers_builder = build_gpu_graph(&test, 1, min_cpu_linked_points_count, true);
+
+        compare_graph_layers_builders(&test.graph_layers_builder, &graph_layers_builder);
+    }
+
+    #[test]
+    fn test_gpu_hnsw_quality_exact() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let num_vectors = 1024;
+        let dim = 64;
+        let m = 8;
+        let m0 = 16;
+        let ef = 32;
+        let groups_count = 4;
+        let searches_count = 20;
+        let top = 10;
+        let min_cpu_linked_points_count = 64;
+
+        let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count);
+        let graph_layers_builder =
+            build_gpu_graph(&test, groups_count, min_cpu_linked_points_count, true);
+
+        check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.8)
+    }
+
+    #[test]
+    fn test_gpu_hnsw_quality() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let num_vectors = 1024;
+        let dim = 64;
+        let m = 8;
+        let m0 = 16;
+        let ef = 32;
+        let groups_count = 4;
+        let searches_count = 20;
+        let top = 10;
+        let min_cpu_linked_points_count = 64;
+
+        let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count);
+        let graph_layers_builder =
+            build_gpu_graph(&test, groups_count, min_cpu_linked_points_count, false);
+
+        check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.8)
+    }
+}

lib/segment/src/index/hnsw_index/gpu/gpu_level_builder.rs

@@ -0,0 +1,211 @@
+use std::sync::atomic::Ordering;
+
+use common::types::PointOffsetType;
+
+use super::batched_points::{Batch, BatchedPoints};
+use super::gpu_search_context::GpuSearchContext;
+use crate::common::operation_error::OperationResult;
+use crate::index::hnsw_index::gpu::gpu_search_context::GpuRequest;
+
+// Build level on GPU
+pub fn build_level_on_gpu(
+    gpu_search_context: &mut GpuSearchContext,
+    batched_points: &BatchedPoints,
+    skip_count: usize,
+    level: usize,
+    processed_points_callback: impl Fn(usize),
+) -> OperationResult<()> {
+    let mut prev_batch = None;
+
+    // skip_count points are already processed
+    processed_points_callback(skip_count);
+
+    for batch in batched_points.iter_batches(skip_count) {
+        if level > batch.level {
+            gpu_batched_update_entries(gpu_search_context, &batch, prev_batch.as_ref())?;
+        } else {
+            gpu_batched_insert(gpu_search_context, &batch, prev_batch.as_ref())?;
+        }
+
+        // Prev batch entries are updated, mark them as processed
+        if let Some(prev_batch) = prev_batch {
+            processed_points_callback(prev_batch.end_index);
+        }
+
+        prev_batch = Some(batch);
+    }
+
+    if let Some(prev_batch) = prev_batch {
+        let new_entries = gpu_search_context.download_responses(prev_batch.points.len())?;
+        gpu_batched_apply_entries(&prev_batch, new_entries);
+        processed_points_callback(prev_batch.end_index);
+    } else {
+        processed_points_callback(0);
+    }
+
+    Ok(())
+}
+
+fn gpu_batched_update_entries(
+    gpu_search_context: &mut GpuSearchContext,
+    batch: &Batch,
+    prev_batch: Option<&Batch>,
+) -> OperationResult<()> {
+    let mut requests = Vec::with_capacity(batch.points.len());
+    for linking_point in batch.points {
+        requests.push(GpuRequest {
+            id: linking_point.point_id,
+            entry: linking_point.entry.load(Ordering::Relaxed),
+        })
+    }
+
+    let prev_batch_len = prev_batch
+        .map(|prev_batch| prev_batch.points.len())
+        .unwrap_or(0);
+
+    let new_entries = gpu_search_context.greedy_search(&requests, prev_batch_len)?;
+
+    if let Some(prev_batch) = prev_batch {
+        gpu_batched_apply_entries(prev_batch, new_entries);
+    }
+    Ok(())
+}
+
+fn gpu_batched_insert(
+    gpu_search_context: &mut GpuSearchContext,
+    batch: &Batch,
+    prev_batch: Option<&Batch>,
+) -> OperationResult<()> {
+    let mut requests = Vec::with_capacity(batch.points.len());
+    for linking_point in batch.points {
+        requests.push(GpuRequest {
+            id: linking_point.point_id,
+            entry: linking_point.entry.load(Ordering::Relaxed),
+        })
+    }
+
+    let prev_batch_len = prev_batch
+        .map(|prev_batch| prev_batch.points.len())
+        .unwrap_or(0);
+
+    let new_entries = gpu_search_context.run_insert_vector(&requests, prev_batch_len)?;
+
+    if let Some(prev_batch) = prev_batch {
+        gpu_batched_apply_entries(prev_batch, new_entries);
+    }
+    Ok(())
+}
+
+fn gpu_batched_apply_entries(batch: &Batch, new_entries: Vec<PointOffsetType>) {
+    assert_eq!(batch.points.len(), new_entries.len());
+    for (linking_point, new_entry) in batch.points.iter().zip(new_entries) {
+        linking_point.entry.store(new_entry, Ordering::Relaxed);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::borrow::Borrow;
+
+    use common::types::PointOffsetType;
+
+    use super::*;
+    use crate::index::hnsw_index::gpu::batched_points::BatchedPoints;
+    use crate::index::hnsw_index::gpu::create_graph_layers_builder;
+    use crate::index::hnsw_index::gpu::tests::{
+        check_graph_layers_builders_quality, compare_graph_layers_builders,
+        create_gpu_graph_test_data, GpuGraphTestData,
+    };
+    use crate::index::hnsw_index::graph_layers::GraphLayersBase;
+    use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder;
+
+    fn build_gpu_graph(test: &GpuGraphTestData, groups_count: usize) -> GraphLayersBuilder {
+        let num_vectors = test.graph_layers_builder.links_layers.len();
+        let m = test.graph_layers_builder.m;
+        let m0 = test.graph_layers_builder.m0;
+        let ef = test.graph_layers_builder.ef_construct;
+
+        let batched_points = BatchedPoints::new(
+            |point_id| test.graph_layers_builder.get_point_level(point_id),
+            (0..num_vectors as PointOffsetType).collect(),
+            groups_count,
+        )
+        .unwrap();
+
+        let graph_layers_builder =
+            create_graph_layers_builder(&batched_points, num_vectors, m, m0, ef, 1).unwrap();
+
+        let debug_messenger = gpu::PanicIfErrorMessenger {};
+        let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap();
+        let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap();
+
+        let mut gpu_search_context = GpuSearchContext::new(
+            device,
+            groups_count,
+            &test.vector_storage.borrow(),
+            None,
+            m,
+            m0,
+            ef,
+            num_vectors,
+            false,
+            true,
+            &false.into(),
+        )
+        .unwrap();
+
+        for level in (0..batched_points.levels_count).rev() {
+            let level_m = graph_layers_builder.get_m(level);
+            gpu_search_context.clear(level_m).unwrap();
+
+            build_level_on_gpu(&mut gpu_search_context, &batched_points, 0, level, |_| {}).unwrap();
+
+            gpu_search_context
+                .download_links(level, &graph_layers_builder)
+                .unwrap();
+        }
+
+        graph_layers_builder
+    }
+
+    #[test]
+    fn test_gpu_hnsw_level_equivalency() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let num_vectors = 1024;
+        let dim = 64;
+        let m = 8;
+        let m0 = 16;
+        let ef = 32;
+
+        let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, 0);
+        let graph_layers_builder = build_gpu_graph(&test, 1);
+
+        compare_graph_layers_builders(&test.graph_layers_builder, &graph_layers_builder);
+    }
+
+    #[test]
+    fn test_gpu_hnsw_level_quality() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let num_vectors = 1024;
+        let dim = 64;
+        let m = 8;
+        let m0 = 16;
+        let ef = 32;
+        let groups_count = 4;
+        let searches_count = 20;
+        let top = 10;
+
+        let test = create_gpu_graph_test_data(num_vectors, dim, m, m0, ef, searches_count);
+        let graph_layers_builder = build_gpu_graph(&test, groups_count);
+
+        check_graph_layers_builders_quality(graph_layers_builder, test, top, ef, 0.8)
+    }
+}

lib/segment/src/index/hnsw_index/gpu/gpu_links.rs

@@ -0,0 +1,298 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use common::types::PointOffsetType;
+
+use super::shader_builder::ShaderBuilderParameters;
+use crate::common::operation_error::{OperationError, OperationResult};
+use crate::index::hnsw_index::gpu::GPU_TIMEOUT;
+use crate::index::hnsw_index::graph_layers::GraphLayersBase;
+use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder;
+
+#[repr(C)]
+struct GpuLinksParamsBuffer {
+    m: u32,
+    links_capacity: u32,
+}
+
+pub struct GpuLinks {
+    pub m: usize,
+    pub links_capacity: usize,
+    pub points_count: usize,
+    pub max_patched_points: usize,
+    pub device: Arc<gpu::Device>,
+    pub links_buffer: Arc<gpu::Buffer>,
+    pub params_buffer: Arc<gpu::Buffer>,
+    pub patch_buffer: Arc<gpu::Buffer>,
+    pub patched_points: Vec<(PointOffsetType, usize)>,
+    pub descriptor_set_layout: Arc<gpu::DescriptorSetLayout>,
+    pub descriptor_set: Arc<gpu::DescriptorSet>,
+}
+
+impl ShaderBuilderParameters for GpuLinks {
+    fn shader_includes(&self) -> HashMap<String, String> {
+        HashMap::from([
+            (
+                "iterators.comp".to_string(),
+                include_str!("shaders/iterators.comp").to_string(),
+            ),
+            (
+                "links.comp".to_string(),
+                include_str!("shaders/links.comp").to_string(),
+            ),
+        ])
+    }
+
+    fn shader_defines(&self) -> HashMap<String, Option<String>> {
+        let mut defines = HashMap::new();
+        defines.insert(
+            "LINKS_CAPACITY".to_owned(),
+            Some(self.links_capacity.to_string()),
+        );
+        defines
+    }
+}
+
+impl GpuLinks {
+    pub fn new(
+        device: Arc<gpu::Device>,
+        m: usize,
+        links_capacity: usize,
+        points_count: usize,
+        max_patched_points: usize,
+    ) -> gpu::GpuResult<Self> {
+        let links_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Links buffer",
+            gpu::BufferType::Storage,
+            points_count * (links_capacity + 1) * std::mem::size_of::<PointOffsetType>(),
+        )?;
+        let params_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Links params buffer",
+            gpu::BufferType::Uniform,
+            std::mem::size_of::<GpuLinksParamsBuffer>(),
+        )?;
+        let links_patch_capacity =
+            max_patched_points * (links_capacity + 1) * std::mem::size_of::<PointOffsetType>();
+        let patch_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Links patch buffer",
+            gpu::BufferType::CpuToGpu,
+            links_patch_capacity + std::mem::size_of::<GpuLinksParamsBuffer>(),
+        )?;
+
+        let params = GpuLinksParamsBuffer {
+            m: m as u32,
+            links_capacity: links_capacity as u32,
+        };
+        patch_buffer.upload(&params, 0)?;
+
+        let mut upload_context = gpu::Context::new(device.clone())?;
+        upload_context.copy_gpu_buffer(
+            patch_buffer.clone(),
+            params_buffer.clone(),
+            0,
+            0,
+            std::mem::size_of::<GpuLinksParamsBuffer>(),
+        )?;
+        upload_context.clear_buffer(links_buffer.clone())?;
+        upload_context.run()?;
+        upload_context.wait_finish(GPU_TIMEOUT)?;
+
+        let descriptor_set_layout = gpu::DescriptorSetLayout::builder()
+            .add_uniform_buffer(0)
+            .add_storage_buffer(1)
+            .build(device.clone())?;
+
+        let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone())
+            .add_uniform_buffer(0, params_buffer.clone())
+            .add_storage_buffer(1, links_buffer.clone())
+            .build()?;
+
+        Ok(Self {
+            m,
+            links_capacity,
+            points_count,
+            max_patched_points,
+            device,
+            links_buffer,
+            params_buffer,
+            patch_buffer,
+            patched_points: vec![],
+            descriptor_set_layout,
+            descriptor_set,
+        })
+    }
+
+    pub fn update_params(&mut self, context: &mut gpu::Context, m: usize) -> OperationResult<()> {
+        self.m = m;
+
+        let params = GpuLinksParamsBuffer {
+            m: m as u32,
+            links_capacity: self.links_capacity as u32,
+        };
+        let links_patch_capacity = self.max_patched_points
+            * (self.links_capacity + 1)
+            * std::mem::size_of::<PointOffsetType>();
+        self.patch_buffer.upload(&params, links_patch_capacity)?;
+
+        context.copy_gpu_buffer(
+            self.patch_buffer.clone(),
+            self.params_buffer.clone(),
+            links_patch_capacity,
+            0,
+            std::mem::size_of::<GpuLinksParamsBuffer>(),
+        )?;
+        Ok(())
+    }
+
+    pub fn clear(&mut self, gpu_context: &mut gpu::Context) -> OperationResult<()> {
+        if !self.patched_points.is_empty() {
+            self.patched_points.clear();
+        }
+        gpu_context.clear_buffer(self.links_buffer.clone())?;
+        Ok(())
+    }
+
+    pub fn apply_gpu_patches(&mut self, gpu_context: &mut gpu::Context) -> OperationResult<()> {
+        for (i, &(patched_point_id, patched_links_count)) in self.patched_points.iter().enumerate()
+        {
+            let patch_start_index =
+                i * (self.links_capacity + 1) * std::mem::size_of::<PointOffsetType>();
+            let patch_size = (patched_links_count + 1) * std::mem::size_of::<PointOffsetType>();
+            let links_start_index = patched_point_id as usize
+                * (self.links_capacity + 1)
+                * std::mem::size_of::<PointOffsetType>();
+            gpu_context.copy_gpu_buffer(
+                self.patch_buffer.clone(),
+                self.links_buffer.clone(),
+                patch_start_index,
+                links_start_index,
+                patch_size,
+            )?;
+        }
+        self.patched_points.clear();
+        Ok(())
+    }
+
+    pub fn set_links(
+        &mut self,
+        point_id: PointOffsetType,
+        links: &[PointOffsetType],
+    ) -> OperationResult<()> {
+        if self.patched_points.len() >= self.max_patched_points {
+            return Err(OperationError::service_error("Gpu links patches are full"));
+        }
+
+        let mut patch_start_index = self.patched_points.len()
+            * (self.links_capacity + 1)
+            * std::mem::size_of::<PointOffsetType>();
+        self.patch_buffer
+            .upload(&(links.len() as u32), patch_start_index)?;
+        patch_start_index += std::mem::size_of::<PointOffsetType>();
+        self.patch_buffer.upload_slice(links, patch_start_index)?;
+        self.patched_points.push((point_id, links.len()));
+
+        Ok(())
+    }
+
+    pub fn upload_links(
+        &mut self,
+        level: usize,
+        graph_layers_builder: &GraphLayersBuilder,
+        context: &mut gpu::Context,
+    ) -> OperationResult<()> {
+        self.update_params(context, graph_layers_builder.get_m(level))?;
+        self.clear(context)?;
+        context.run()?;
+        context.wait_finish(GPU_TIMEOUT)?;
+
+        let timer = std::time::Instant::now();
+        let points: Vec<_> = (0..graph_layers_builder.links_layers.len())
+            .filter(|&point_id| {
+                graph_layers_builder.get_point_level(point_id as PointOffsetType) >= level
+            })
+            .filter(|&point_id| {
+                !graph_layers_builder.links_layers[point_id][level]
+                    .read()
+                    .is_empty()
+            })
+            .collect();
+
+        for points_slice in points.chunks(self.max_patched_points) {
+            for &point_id in points_slice {
+                let links = graph_layers_builder.links_layers[point_id][level].read();
+                self.set_links(point_id as PointOffsetType, &links)?;
+            }
+            self.apply_gpu_patches(context)?;
+            context.run()?;
+            context.wait_finish(GPU_TIMEOUT)?;
+        }
+
+        log::trace!("Upload links on level {level} time: {:?}", timer.elapsed());
+        Ok(())
+    }
+
+    pub fn download_links(
+        &mut self,
+        level: usize,
+        graph_layers_builder: &GraphLayersBuilder,
+        context: &mut gpu::Context,
+    ) -> OperationResult<()> {
+        let timer = std::time::Instant::now();
+
+        let links_patch_capacity = self.max_patched_points
+            * (self.links_capacity + 1)
+            * std::mem::size_of::<PointOffsetType>();
+        let download_buffer = gpu::Buffer::new(
+            self.device.clone(),
+            "Download links staging buffer",
+            gpu::BufferType::GpuToCpu,
+            links_patch_capacity,
+        )
+        .unwrap();
+
+        let points = (0..graph_layers_builder.links_layers.len() as PointOffsetType)
+            .filter(|&point_id| graph_layers_builder.get_point_level(point_id) >= level)
+            .collect::<Vec<_>>();
+
+        for chunk_index in 0..points.len().div_ceil(self.max_patched_points) {
+            let start = chunk_index * self.max_patched_points;
+            let end = (start + self.max_patched_points).min(points.len());
+            let chunk_size = end - start;
+            for (i, &point_id) in points[start..end].iter().enumerate() {
+                let links_size = (self.links_capacity + 1) * std::mem::size_of::<PointOffsetType>();
+                context.copy_gpu_buffer(
+                    self.links_buffer.clone(),
+                    download_buffer.clone(),
+                    point_id as usize * links_size,
+                    i * links_size,
+                    links_size,
+                )?;
+            }
+            context.run()?;
+            context.wait_finish(GPU_TIMEOUT)?;
+
+            let mut links =
+                vec![PointOffsetType::default(); chunk_size * (self.links_capacity + 1)];
+            download_buffer.download_slice(&mut links, 0)?;
+
+            for (index, chunk) in links.chunks(self.links_capacity + 1).enumerate() {
+                let point_id = points[start + index] as usize;
+                let links_count = chunk[0] as usize;
+                let links = &chunk[1..=links_count];
+                let mut dst = graph_layers_builder.links_layers[point_id][level].write();
+                dst.clear();
+                dst.extend_from_slice(links);
+            }
+        }
+
+        log::trace!(
+            "Download links for level {} in time {:?}",
+            level,
+            timer.elapsed()
+        );
+        Ok(())
+    }
+}

lib/segment/src/index/hnsw_index/gpu/gpu_nearest_heap.rs

@@ -0,0 +1,259 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use super::shader_builder::ShaderBuilderParameters;
+
+pub struct GpuNearestHeap {
+    pub ef: usize,
+    pub capacity: usize,
+    pub device: Arc<gpu::Device>,
+}
+
+impl ShaderBuilderParameters for GpuNearestHeap {
+    fn shader_includes(&self) -> HashMap<String, String> {
+        HashMap::from([(
+            "nearest_heap.comp".to_string(),
+            include_str!("shaders/nearest_heap.comp").to_string(),
+        )])
+    }
+
+    fn shader_defines(&self) -> HashMap<String, Option<String>> {
+        let mut defines = HashMap::new();
+        defines.insert("NEAREST_HEAP_EF".to_owned(), Some(self.ef.to_string()));
+        defines.insert(
+            "NEAREST_HEAP_CAPACITY".to_owned(),
+            Some(self.capacity.to_string()),
+        );
+        defines
+    }
+}
+
+impl GpuNearestHeap {
+    pub fn new(device: Arc<gpu::Device>, ef: usize, capacity: usize) -> gpu::GpuResult<Self> {
+        assert!(capacity >= ef);
+        let capacity = capacity.div_ceil(device.subgroup_size()) * device.subgroup_size();
+        Ok(Self {
+            ef,
+            capacity,
+            device,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use common::fixed_length_priority_queue::FixedLengthPriorityQueue;
+    use common::types::{PointOffsetType, ScoredPointOffset};
+    use rand::rngs::StdRng;
+    use rand::{Rng, SeedableRng};
+
+    use super::*;
+    use crate::index::hnsw_index::gpu::shader_builder::ShaderBuilder;
+    use crate::index::hnsw_index::gpu::GPU_TIMEOUT;
+
+    #[repr(C)]
+    struct TestParams {
+        input_counts: u32,
+    }
+
+    #[test]
+    fn test_gpu_nearest_heap() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(log::LevelFilter::Trace)
+            .try_init();
+
+        let ef = 100;
+        let points_count = 1024;
+        let groups_count = 8;
+        let inputs_count = points_count;
+
+        let mut rng = StdRng::seed_from_u64(42);
+        let inputs_data: Vec<ScoredPointOffset> = (0..inputs_count * groups_count)
+            .map(|i| ScoredPointOffset {
+                idx: i as PointOffsetType,
+                score: rng.gen_range(-1.0..1.0),
+            })
+            .collect();
+
+        let debug_messenger = gpu::PanicIfErrorMessenger {};
+        let instance = gpu::Instance::new(Some(&debug_messenger), None, false).unwrap();
+        let device = gpu::Device::new(instance.clone(), &instance.physical_devices()[0]).unwrap();
+
+        let gpu_nearest_heap = GpuNearestHeap::new(device.clone(), ef, ef).unwrap();
+
+        let shader = ShaderBuilder::new(device.clone())
+            .with_shader_code(include_str!("shaders/tests/test_nearest_heap.comp"))
+            .with_parameters(&gpu_nearest_heap)
+            .build()
+            .unwrap();
+
+        let input_points_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Nearest heap input points buffer",
+            gpu::BufferType::Storage,
+            inputs_count * groups_count * std::mem::size_of::<ScoredPointOffset>(),
+        )
+        .unwrap();
+
+        let upload_staging_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Nearest heap upload staging buffer",
+            gpu::BufferType::CpuToGpu,
+            inputs_count * groups_count * std::mem::size_of::<ScoredPointOffset>(),
+        )
+        .unwrap();
+        upload_staging_buffer.upload_slice(&inputs_data, 0).unwrap();
+
+        let mut context = gpu::Context::new(device.clone()).unwrap();
+        context
+            .copy_gpu_buffer(
+                upload_staging_buffer.clone(),
+                input_points_buffer.clone(),
+                0,
+                0,
+                input_points_buffer.size(),
+            )
+            .unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+
+        let test_params_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Nearest heap test params buffer",
+            gpu::BufferType::Uniform,
+            std::mem::size_of::<TestParams>(),
+        )
+        .unwrap();
+        upload_staging_buffer
+            .upload(
+                &TestParams {
+                    input_counts: inputs_count as u32,
+                },
+                0,
+            )
+            .unwrap();
+        context
+            .copy_gpu_buffer(
+                upload_staging_buffer,
+                test_params_buffer.clone(),
+                0,
+                0,
+                test_params_buffer.size(),
+            )
+            .unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+
+        let scores_output_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Nearest heap scores output buffer",
+            gpu::BufferType::Storage,
+            inputs_count * groups_count * std::mem::size_of::<f32>(),
+        )
+        .unwrap();
+        let sorted_output_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Nearest heap sorted output buffer",
+            gpu::BufferType::Storage,
+            ef * groups_count * std::mem::size_of::<PointOffsetType>(),
+        )
+        .unwrap();
+
+        let descriptor_set_layout = gpu::DescriptorSetLayout::builder()
+            .add_uniform_buffer(0)
+            .add_storage_buffer(1)
+            .add_storage_buffer(2)
+            .add_storage_buffer(3)
+            .build(device.clone())
+            .unwrap();
+
+        let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone())
+            .add_uniform_buffer(0, test_params_buffer.clone())
+            .add_storage_buffer(1, input_points_buffer.clone())
+            .add_storage_buffer(2, scores_output_buffer.clone())
+            .add_storage_buffer(3, sorted_output_buffer.clone())
+            .build()
+            .unwrap();
+
+        let pipeline = gpu::Pipeline::builder()
+            .add_descriptor_set_layout(0, descriptor_set_layout.clone())
+            .add_shader(shader.clone())
+            .build(device.clone())
+            .unwrap();
+
+        context
+            .bind_pipeline(pipeline, &[descriptor_set.clone()])
+            .unwrap();
+        context.dispatch(groups_count, 1, 1).unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+
+        let download_staging_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Nearest heap download staging buffer",
+            gpu::BufferType::GpuToCpu,
+            std::cmp::max(scores_output_buffer.size(), sorted_output_buffer.size()),
+        )
+        .unwrap();
+        context
+            .copy_gpu_buffer(
+                scores_output_buffer.clone(),
+                download_staging_buffer.clone(),
+                0,
+                0,
+                scores_output_buffer.size(),
+            )
+            .unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+        let mut scores_output = vec![0.0; inputs_count * groups_count];
+        download_staging_buffer
+            .download_slice(&mut scores_output, 0)
+            .unwrap();
+
+        let mut scores_output_cpu = vec![0.0; inputs_count * groups_count];
+        let mut sorted_output_cpu = vec![PointOffsetType::default(); ef * groups_count];
+        for group in 0..groups_count {
+            let mut queue = FixedLengthPriorityQueue::<ScoredPointOffset>::new(ef);
+            for i in 0..inputs_count {
+                let scored_point = inputs_data[group * inputs_count + i];
+                queue.push(scored_point);
+                scores_output_cpu[group * inputs_count + i] = queue.top().unwrap().score;
+            }
+            let sorted = queue.into_vec();
+            for i in 0..ef {
+                sorted_output_cpu[group * ef + i] = sorted[i].idx;
+            }
+        }
+
+        let mut nearest_gpu: Vec<PointOffsetType> =
+            vec![Default::default(); gpu_nearest_heap.ef * groups_count];
+        context
+            .copy_gpu_buffer(
+                sorted_output_buffer.clone(),
+                download_staging_buffer.clone(),
+                0,
+                0,
+                nearest_gpu.len() * std::mem::size_of::<PointOffsetType>(),
+            )
+            .unwrap();
+        context.run().unwrap();
+        context.wait_finish(GPU_TIMEOUT).unwrap();
+        download_staging_buffer
+            .download_slice(nearest_gpu.as_mut_slice(), 0)
+            .unwrap();
+
+        let mut sorted_output_gpu = Vec::new();
+        for group in 0..groups_count {
+            let mut nearest_group = Vec::new();
+            for i in 0..ef {
+                nearest_group.push(nearest_gpu[group * gpu_nearest_heap.ef + i]);
+            }
+            sorted_output_gpu.extend(nearest_group);
+        }
+
+        assert_eq!(scores_output, scores_output_cpu);
+        assert_eq!(sorted_output_gpu, sorted_output_cpu);
+    }
+}

lib/segment/src/index/hnsw_index/gpu/gpu_visited_flags.rs

@@ -0,0 +1,124 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use super::shader_builder::ShaderBuilderParameters;
+use super::GPU_TIMEOUT;
+use crate::common::operation_error::OperationResult;
+
+#[repr(C)]
+pub struct GpuVisitedFlagsParamsBuffer {
+    pub generation: u32,
+}
+
+pub struct GpuVisitedFlags {
+    pub device: Arc<gpu::Device>,
+    pub params: GpuVisitedFlagsParamsBuffer,
+    pub params_buffer: Arc<gpu::Buffer>,
+    pub params_staging_buffer: Arc<gpu::Buffer>,
+    pub visited_flags_buffer: Arc<gpu::Buffer>,
+    pub descriptor_set_layout: Arc<gpu::DescriptorSetLayout>,
+    pub descriptor_set: Arc<gpu::DescriptorSet>,
+    pub capacity: usize,
+}
+
+impl ShaderBuilderParameters for GpuVisitedFlags {
+    fn shader_includes(&self) -> HashMap<String, String> {
+        HashMap::from([(
+            "visited_flags.comp".to_string(),
+            include_str!("shaders/visited_flags.comp").to_string(),
+        )])
+    }
+
+    fn shader_defines(&self) -> HashMap<String, Option<String>> {
+        let mut defines = HashMap::new();
+        defines.insert(
+            "VISITED_FLAGS_CAPACITY".to_owned(),
+            Some(self.capacity.to_string()),
+        );
+        defines
+    }
+}
+
+impl GpuVisitedFlags {
+    pub fn new(
+        device: Arc<gpu::Device>,
+        groups_count: usize,
+        points_count: usize,
+    ) -> OperationResult<Self> {
+        let alignment = std::mem::size_of::<u32>();
+        let points_count = points_count.div_ceil(alignment) * alignment;
+
+        let params_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Visited flags params buffer",
+            gpu::BufferType::Uniform,
+            std::mem::size_of::<GpuVisitedFlagsParamsBuffer>(),
+        )?;
+        let params_staging_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Visited flags params staging buffer",
+            gpu::BufferType::CpuToGpu,
+            std::mem::size_of::<GpuVisitedFlagsParamsBuffer>(),
+        )?;
+        let visited_flags_buffer = gpu::Buffer::new(
+            device.clone(),
+            "Visited flags buffer",
+            gpu::BufferType::Storage,
+            groups_count * points_count * std::mem::size_of::<u8>(),
+        )?;
+
+        let params = GpuVisitedFlagsParamsBuffer { generation: 1 };
+        params_staging_buffer.upload(&params, 0)?;
+
+        let mut upload_context = gpu::Context::new(device.clone())?;
+        upload_context.copy_gpu_buffer(
+            params_staging_buffer.clone(),
+            params_buffer.clone(),
+            0,
+            0,
+            std::mem::size_of::<GpuVisitedFlagsParamsBuffer>(),
+        )?;
+        upload_context.run()?;
+        upload_context.wait_finish(GPU_TIMEOUT)?;
+
+        let descriptor_set_layout = gpu::DescriptorSetLayout::builder()
+            .add_uniform_buffer(0)
+            .add_storage_buffer(1)
+            .build(device.clone())?;
+
+        let descriptor_set = gpu::DescriptorSet::builder(descriptor_set_layout.clone())
+            .add_uniform_buffer(0, params_buffer.clone())
+            .add_storage_buffer(1, visited_flags_buffer.clone())
+            .build()?;
+
+        Ok(Self {
+            device,
+            params,
+            params_buffer,
+            params_staging_buffer,
+            visited_flags_buffer,
+            descriptor_set_layout,
+            descriptor_set,
+            capacity: points_count,
+        })
+    }
+
+    pub fn clear(&mut self, gpu_context: &mut gpu::Context) -> OperationResult<()> {
+        if self.params.generation == 255 {
+            self.params.generation = 1;
+            gpu_context.clear_buffer(self.visited_flags_buffer.clone())?;
+        } else {
+            self.params.generation += 1;
+        }
+
+        self.params_staging_buffer.upload(&self.params, 0)?;
+        gpu_context.copy_gpu_buffer(
+            self.params_staging_buffer.clone(),
+            self.params_buffer.clone(),
+            0,
+            0,
+            self.params_buffer.size(),
+        )?;
+        Ok(())
+    }
+}

lib/segment/src/index/hnsw_index/gpu/mod.rs

@@ -1,8 +1,353 @@
+pub mod batched_points;
+pub mod cpu_level_builder;
+pub mod devices_manager;
+pub mod gpu_candidates_heap;
+pub mod gpu_graph_builder;
+pub mod gpu_level_builder;
+pub mod gpu_links;
+pub mod gpu_nearest_heap;
+pub mod gpu_search_context;
 pub mod gpu_vector_storage;
+pub mod gpu_visited_flags;
 pub mod shader_builder;
 
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::Arc;
+
+use batched_points::BatchedPoints;
+use bitvec::vec::BitVec;
+use devices_manager::DevicesMaganer;
+use lazy_static::lazy_static;
+use parking_lot::{Mutex, RwLock};
+
+use super::graph_layers_builder::GraphLayersBuilder;
+use crate::common::operation_error::{OperationError, OperationResult};
+
+lazy_static! {
+    static ref GPU_INSTANCE: OperationResult<Arc<gpu::Instance>> = create_gpu_instance();
+    static ref GPU_DEVICE: OperationResult<Arc<gpu::Device>> = create_gpu_device();
+    static ref GPU_DEVICE_FILER: Mutex<String> = Mutex::new("".to_string());
+    pub static ref GPU_DEVICES_MANAGER: OperationResult<DevicesMaganer> = init_devices_manager();
+}
+
 /// Each GPU operation has a timeout by Vulkan API specification.
 /// Choose large enough timeout.
 /// We cannot use too small timeout and check stopper in the loop because
 /// GPU resources should be alive while GPU operation is in progress.
 static GPU_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
+
+static GPU_INDEXING: AtomicBool = AtomicBool::new(false);
+static GPU_WAIT_FREE: AtomicBool = AtomicBool::new(true);
+static GPU_FORCE_HALF_PRECISION: AtomicBool = AtomicBool::new(false);
+static GPU_MAX_GROUPS: AtomicUsize = AtomicUsize::new(GPU_MAX_GROUPS_COUNT_DEFAULT);
+pub const GPU_MAX_GROUPS_COUNT_DEFAULT: usize = 512;
+static GPU_MIN_POINTS_COUNT: AtomicUsize = AtomicUsize::new(10_000);
+static GPU_DEVICE_START_INDEX: AtomicUsize = AtomicUsize::new(0);
+static GPU_DEVICES_COUNT: AtomicUsize = AtomicUsize::new(usize::MAX);
+static GPU_PARALLEL_INDEXES: AtomicUsize = AtomicUsize::new(0);
+
+fn create_gpu_instance() -> OperationResult<Arc<gpu::Instance>> {
+    Ok(gpu::Instance::new(None, None, false)?)
+}
+
+fn init_devices_manager() -> OperationResult<DevicesMaganer> {
+    let instance = GPU_INSTANCE.clone()?;
+    let filter = GPU_DEVICE_FILER.lock().clone();
+    // TODO(gpu): add start index and count
+    let devices_manager = DevicesMaganer::new(
+        instance,
+        &filter,
+        GPU_DEVICE_START_INDEX.load(Ordering::Relaxed),
+        GPU_DEVICES_COUNT.load(Ordering::Relaxed),
+        GPU_WAIT_FREE.load(Ordering::Relaxed),
+        GPU_PARALLEL_INDEXES.load(Ordering::Relaxed),
+    )?;
+    Ok(devices_manager)
+}
+
+fn create_gpu_device() -> OperationResult<Arc<gpu::Device>> {
+    let instance = GPU_INSTANCE.clone()?;
+    let filter = GPU_DEVICE_FILER.lock().clone();
+
+    for physical_device in instance.physical_devices() {
+        let device_name = physical_device.name.to_owned();
+        if filter.is_empty() || device_name.contains(&filter) {
+            log::info!("Found GPU device: {}", device_name);
+            return Ok(gpu::Device::new(instance.clone(), physical_device)?);
+        }
+    }
+
+    Err(OperationError::service_error("No GPU device found"))
+}
+
+pub fn set_gpu_indexing(gpu_indexing: bool) {
+    GPU_INDEXING.store(gpu_indexing, Ordering::Relaxed);
+}
+
+pub fn get_gpu_indexing() -> bool {
+    GPU_INDEXING.load(Ordering::Relaxed)
+}
+
+pub fn set_wait_free(wait_free: bool) {
+    GPU_WAIT_FREE.store(wait_free, Ordering::Relaxed);
+}
+
+pub fn set_gpu_force_half_precision(force_half_precision: bool) {
+    GPU_FORCE_HALF_PRECISION.store(force_half_precision, Ordering::Relaxed);
+}
+
+pub fn get_gpu_force_half_precision() -> bool {
+    GPU_FORCE_HALF_PRECISION.load(Ordering::Relaxed)
+}
+
+pub fn set_gpu_max_groups(max_groups: Option<usize>) {
+    if let Some(max_groups) = max_groups {
+        GPU_MAX_GROUPS.store(max_groups, Ordering::Relaxed);
+    }
+}
+
+pub fn get_gpu_max_groups() -> usize {
+    GPU_MAX_GROUPS.load(Ordering::Relaxed)
+}
+
+pub fn set_device_index(device_index: Option<usize>) {
+    if let Some(device_index) = device_index {
+        GPU_DEVICE_START_INDEX.store(device_index, Ordering::Relaxed);
+    }
+}
+
+pub fn set_devices_count(devices_count: Option<usize>) {
+    if let Some(devices_count) = devices_count {
+        GPU_DEVICES_COUNT.store(devices_count, Ordering::Relaxed);
+    }
+}
+
+pub fn set_device_filter(device_filter: &str) {
+    let mut filter = GPU_DEVICE_FILER.lock();
+    *filter = device_filter.to_string();
+}
+
+pub fn set_gpu_parallel_indexes(parallel_indexes: usize) {
+    GPU_PARALLEL_INDEXES.store(parallel_indexes, Ordering::Relaxed);
+}
+
+pub fn get_gpu_min_points_count() -> usize {
+    GPU_MIN_POINTS_COUNT.load(Ordering::Relaxed)
+}
+
+pub fn create_graph_layers_builder(
+    batched_points: &BatchedPoints,
+    num_vectors: usize,
+    m: usize,
+    m0: usize,
+    ef: usize,
+    entry_points_num: usize,
+) -> OperationResult<GraphLayersBuilder> {
+    // create graph layers builder
+    let mut graph_layers_builder =
+        GraphLayersBuilder::new(num_vectors, m, m0, ef, entry_points_num, true);
+
+    // mark all vectors as ready
+    graph_layers_builder.ready_list = RwLock::new(BitVec::repeat(true, num_vectors));
+
+    // set first entry point
+    graph_layers_builder.set_levels(
+        batched_points.first_point_id,
+        batched_points.levels_count - 1,
+    );
+    graph_layers_builder.entry_points.lock().new_point(
+        batched_points.first_point_id,
+        batched_points.levels_count - 1,
+        |_| true,
+    );
+
+    // set levels
+    for linking_point in &batched_points.points {
+        graph_layers_builder.set_levels(linking_point.point_id, linking_point.level);
+        let mut entry_points = graph_layers_builder.entry_points.lock();
+        entry_points.new_point(linking_point.point_id, linking_point.level, |_| true);
+    }
+
+    Ok(graph_layers_builder)
+}
+
+#[cfg(test)]
+mod tests {
+    use ahash::HashSet;
+    use common::types::PointOffsetType;
+    use rand::rngs::StdRng;
+    use rand::SeedableRng;
+    use tempfile::TempDir;
+
+    use super::batched_points::BatchedPoints;
+    use crate::common::rocksdb_wrapper::{open_db, DB_VECTOR_CF};
+    use crate::data_types::vectors::DenseVector;
+    use crate::fixtures::index_fixtures::{FakeFilterContext, TestRawScorerProducer};
+    use crate::fixtures::payload_fixtures::random_vector;
+    use crate::index::hnsw_index::graph_layers::GraphLayers;
+    use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder;
+    use crate::index::hnsw_index::graph_links::GraphLinksRam;
+    use crate::index::hnsw_index::point_scorer::FilteredScorer;
+    use crate::spaces::simple::CosineMetric;
+    use crate::types::Distance;
+    use crate::vector_storage::chunked_vector_storage::VectorOffsetType;
+    use crate::vector_storage::dense::simple_dense_vector_storage::open_simple_dense_vector_storage;
+    use crate::vector_storage::{VectorStorage, VectorStorageEnum};
+
+    #[allow(dead_code)]
+    pub struct GpuGraphTestData {
+        pub dir: TempDir,
+        pub vector_storage: VectorStorageEnum,
+        pub vector_holder: TestRawScorerProducer<CosineMetric>,
+        pub graph_layers_builder: GraphLayersBuilder,
+        pub search_vectors: Vec<DenseVector>,
+    }
+
+    pub fn create_gpu_graph_test_data(
+        num_vectors: usize,
+        dim: usize,
+        m: usize,
+        m0: usize,
+        ef: usize,
+        search_counts: usize,
+    ) -> GpuGraphTestData {
+        // Generate random vectors
+        let mut rng = StdRng::seed_from_u64(42);
+        let vector_holder = TestRawScorerProducer::<CosineMetric>::new(dim, num_vectors, &mut rng);
+
+        // upload vectors to storage
+        let dir = tempfile::Builder::new().prefix("db_dir").tempdir().unwrap();
+        let db = open_db(dir.path(), &[DB_VECTOR_CF]).unwrap();
+        let mut storage = open_simple_dense_vector_storage(
+            db,
+            DB_VECTOR_CF,
+            dim,
+            Distance::Cosine,
+            &false.into(),
+        )
+        .unwrap();
+        for idx in 0..num_vectors {
+            let v = vector_holder.get_vector(idx as PointOffsetType);
+            storage
+                .insert_vector(idx as PointOffsetType, v.as_vec_ref())
+                .unwrap();
+        }
+
+        // Build HNSW index
+        let mut graph_layers_builder = GraphLayersBuilder::new(num_vectors, m, m0, ef, 1, true);
+        for idx in 0..(num_vectors as PointOffsetType) {
+            let level = graph_layers_builder.get_random_layer(&mut rng);
+            graph_layers_builder.set_levels(idx, level);
+        }
+
+        let mut ids: Vec<_> = (0..num_vectors as PointOffsetType).collect();
+        BatchedPoints::sort_points_by_level(
+            &|point_id| graph_layers_builder.get_point_level(point_id),
+            &mut ids,
+        );
+
+        for &idx in &ids {
+            let fake_filter_context = FakeFilterContext {};
+            let added_vector = vector_holder.vectors.get(idx as VectorOffsetType).to_vec();
+            let raw_scorer = vector_holder.get_raw_scorer(added_vector.clone()).unwrap();
+            let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context));
+            graph_layers_builder.link_new_point(idx, scorer);
+            raw_scorer.take_hardware_counter().discard_results();
+        }
+
+        let search_vectors = (0..search_counts)
+            .map(|_| random_vector(&mut rng, dim))
+            .collect();
+
+        GpuGraphTestData {
+            dir,
+            vector_storage: storage,
+            vector_holder,
+            graph_layers_builder,
+            search_vectors,
+        }
+    }
+
+    pub fn compare_graph_layers_builders(
+        graph_a: &GraphLayersBuilder,
+        graph_b: &GraphLayersBuilder,
+    ) {
+        assert_eq!(graph_a.links_layers.len(), graph_b.links_layers.len());
+        let num_vectors = graph_a.links_layers.len();
+        for point_id in 0..num_vectors as PointOffsetType {
+            let levels_a = graph_a.get_point_level(point_id);
+            let levels_b = graph_b.get_point_level(point_id);
+            assert_eq!(levels_a, levels_b);
+
+            for level in (0..levels_a + 1).rev() {
+                let links_a = graph_a.links_layers[point_id as usize][level]
+                    .read()
+                    .clone();
+                let links_b = graph_b.links_layers[point_id as usize][level]
+                    .read()
+                    .clone();
+                if links_a != links_b {
+                    log::error!("Wrong links point_id={} at level {}", point_id, level);
+                }
+                assert_eq!(links_a, links_b);
+            }
+        }
+    }
+
+    pub fn check_graph_layers_builders_quality(
+        graph: GraphLayersBuilder,
+        test: GpuGraphTestData,
+        top: usize,
+        ef: usize,
+        accuracy: f32,
+    ) {
+        let graph: GraphLayers<GraphLinksRam> = graph.into_graph_layers(None).unwrap();
+        let ref_graph: GraphLayers<GraphLinksRam> =
+            test.graph_layers_builder.into_graph_layers(None).unwrap();
+
+        let mut total_sames = 0;
+        let total_top = top * test.search_vectors.len();
+        for search_vector in &test.search_vectors {
+            let fake_filter_context = FakeFilterContext {};
+            let raw_scorer = test
+                .vector_holder
+                .get_raw_scorer(search_vector.clone())
+                .unwrap();
+            let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context));
+
+            let search_result_gpu = graph.search(top, ef, scorer, None);
+            raw_scorer.take_hardware_counter().discard_results();
+
+            let fake_filter_context = FakeFilterContext {};
+            let raw_scorer = test
+                .vector_holder
+                .get_raw_scorer(search_vector.clone())
+                .unwrap();
+            let scorer = FilteredScorer::new(raw_scorer.as_ref(), Some(&fake_filter_context));
+
+            let search_result_cpu = ref_graph.search(top, ef, scorer, None);
+            raw_scorer.take_hardware_counter().discard_results();
+
+            let mut gpu_set = HashSet::default();
+            let mut cpu_set = HashSet::default();
+            for (gpu_id, cpu_id) in search_result_gpu.iter().zip(search_result_cpu.iter()) {
+                gpu_set.insert(gpu_id.idx);
+                cpu_set.insert(cpu_id.idx);
+            }
+
+            total_sames += gpu_set.intersection(&cpu_set).count();
+        }
+        log::debug!(
+            "total_sames: {}, total_top: {}, div {}",
+            total_sames,
+            total_top,
+            total_sames as f32 / total_top as f32
+        );
+        assert!(
+            total_sames as f32 >= total_top as f32 * accuracy,
+            "sames: {}, total_top: {}",
+            total_sames,
+            total_top
+        );
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/bheap.comp

@@ -0,0 +1,123 @@
+/******
+Shader binary heap.
+
+Usage example.
+Step 1: defines
+
+#define BHEAP_NAME bheap // provide binary heap name. Name is required for cases when there are many binary heaps in one shader
+float heap_data[8]; // reserve storage for binary heap
+#define BHEAP_DATA heap_data // provide storage name
+#define BHEAP_TYPE float // binary tree element type
+#define BHEAP_CMP(a, b) a > b // cmp function
+#include "bheap.comp" // define all binary tree functions
+
+Step 2: init
+BHEAP(BHEAP_NAME, init)(0.0, 8, 0); // default value, capacity and offset in storage
+
+Step 3: usage
+BHEAP(BHEAP_NAME, capacity) // get capacity
+BHEAP(BHEAP_NAME, offset) // get offset
+BHEAP(BHEAP_NAME, size); // get size
+BHEAP(BHEAP_NAME, top)(); // get top value
+BHEAP(BHEAP_NAME, pop)(); // pop value
+BHEAP(BHEAP_NAME, change_top)(1.0); // pop and push
+BHEAP(BHEAP_NAME, push)(1.0); // push value
+BHEAP(BHEAP_NAME, into_array)(); // convert to sorted array and return array size
+
+******/
+
+#include "common.comp"
+
+#define BHEAP(NAME, a) CONCAT(NAME, a)
+
+uint BHEAP(BHEAP_NAME, offset) = 0;
+uint BHEAP(BHEAP_NAME, capacity) = 0;
+uint BHEAP(BHEAP_NAME, size) = 0;
+
+// swap
+void BHEAP(BHEAP_NAME, swap)(uint i, uint j) {
+    BHEAP_TYPE tmp = BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i];
+    BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + j];
+    BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + j] = tmp;
+}
+
+// compare, call BHEAP_CMP function
+bool BHEAP(BHEAP_NAME, cmp)(uint i, uint j) {
+    return BHEAP_CMP(
+        (BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i]),
+        (BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + j])
+    );
+}
+
+// sift down
+void BHEAP(BHEAP_NAME, sift_down)(uint i) {
+    while (2 * i + 1 < BHEAP(BHEAP_NAME, size)) {
+        uint left = 2 * i + 1;
+        uint right = 2 * i + 2;
+        uint j = left;
+        if (right < BHEAP(BHEAP_NAME, size) && BHEAP(BHEAP_NAME, cmp)(right, left)) {
+            j = right;
+        }
+        if (BHEAP(BHEAP_NAME, cmp)(i, j)) {
+            break;
+        }
+        BHEAP(BHEAP_NAME, swap)(i, j);
+        i = j;
+    }
+}
+
+// push
+void BHEAP(BHEAP_NAME, push)(BHEAP_TYPE score) {
+    uint i = BHEAP(BHEAP_NAME, size);
+    BHEAP(BHEAP_NAME, size)++;
+    BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = score;
+    while (i > 0 && BHEAP(BHEAP_NAME, cmp)(i, (i - 1) / 2)) {
+        BHEAP(BHEAP_NAME, swap)(i, (i - 1) / 2);
+        i = (i - 1) / 2;
+    }
+}
+
+// get top value
+BHEAP_TYPE BHEAP(BHEAP_NAME, top)() {
+    return BHEAP_DATA[BHEAP(BHEAP_NAME, offset)];
+}
+
+// pop value
+BHEAP_TYPE BHEAP(BHEAP_NAME, pop)() {
+    BHEAP_TYPE top_value = BHEAP(BHEAP_NAME, top)();
+    BHEAP(BHEAP_NAME, size)--;
+    BHEAP_DATA[BHEAP(BHEAP_NAME, offset)] = BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + BHEAP(BHEAP_NAME, size)];
+    BHEAP(BHEAP_NAME, sift_down)(0);
+    return top_value;
+}
+
+// replace top value
+void BHEAP(BHEAP_NAME, change_top)(BHEAP_TYPE new_value) {
+    BHEAP_DATA[BHEAP(BHEAP_NAME, offset)] = new_value;
+    BHEAP(BHEAP_NAME, sift_down)(0);
+}
+
+// into_array
+uint BHEAP(BHEAP_NAME, into_array)() {
+    uint array_size = BHEAP(BHEAP_NAME, size);
+    while (BHEAP(BHEAP_NAME, size) > 0) {
+        uint i = BHEAP(BHEAP_NAME, size) - 1;
+        BHEAP_TYPE value = BHEAP(BHEAP_NAME, pop)();
+        BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = value;
+    }
+    return array_size;
+}
+
+// init
+void BHEAP(BHEAP_NAME, init)(BHEAP_TYPE none_value, uint capacity, uint offset) {
+    BHEAP(BHEAP_NAME, offset) = offset;
+    BHEAP(BHEAP_NAME, capacity) = capacity;
+    //for (uint i = 0; i < capacity; i++) {
+    //    BHEAP_DATA[BHEAP(BHEAP_NAME, offset) + i] = none_value;
+    //}
+}
+
+#undef BHEAP_NAME
+#undef BHEAP_DATA
+#undef BHEAP_TYPE
+#undef BHEAP_CMP

lib/segment/src/index/hnsw_index/gpu/shaders/candidates_heap.comp

@@ -0,0 +1,75 @@
+#ifndef CANDIDATES_HEAP_HEADER
+#define CANDIDATES_HEAP_HEADER
+
+#include "common.comp"
+
+shared ScoredPoint candidates[CANDIDATES_HEAP_CAPACITY];
+uint candidates_count;
+
+void candidates_heap_init() {
+    candidates_count = 0;
+    for (uint i = gl_SubgroupInvocationID; i < CANDIDATES_HEAP_CAPACITY; i += SUBGROUP_SIZE) {
+        candidates[i] = ScoredPoint(0, positive_infinity);
+    }
+}
+
+ScoredPoint pop_candidate() {
+    if (candidates_count == 0) {
+        return ScoredPoint(0, 0.0);
+    }
+
+    ScoredPoint max_scored_point = ScoredPoint(0, negative_infinity);
+    uint max_index = 0;
+    for (uint i = gl_SubgroupInvocationID; i < candidates_count; i += SUBGROUP_SIZE) {
+        if (candidates[i].score > max_scored_point.score) {
+            max_scored_point = candidates[i];
+            max_index = i;
+        }
+    }
+
+    candidates_count--;
+
+    float subgroup_largest = subgroupMax(max_scored_point.score);
+    bool is_max_thread = max_scored_point.score == subgroup_largest;
+    uvec4 subgroup_max_threads = subgroupBallot(is_max_thread);
+    uint subgroup_max_thread = subgroupBallotFindLSB(subgroup_max_threads);
+    max_index = subgroupShuffle(max_index, subgroup_max_thread);
+    max_scored_point.score = subgroupShuffle(max_scored_point.score, subgroup_max_thread);
+    max_scored_point.id = subgroupShuffle(max_scored_point.id, subgroup_max_thread);
+
+    if (subgroupElect()) {
+        candidates[max_index] = candidates[candidates_count];
+    }
+    return max_scored_point;
+}
+
+bool push_candidate(ScoredPoint scored_point) {
+    if (candidates_count == CANDIDATES_HEAP_CAPACITY) {
+        ScoredPoint min_scored_point = ScoredPoint(0, positive_infinity);
+        uint min_index = 0;
+        for (uint i = gl_SubgroupInvocationID; i < candidates_count; i += SUBGROUP_SIZE) {
+            if (candidates[i].score < min_scored_point.score) {
+                min_scored_point = candidates[i];
+                min_index = i;
+            }
+        }
+
+        float subgroup_smallest = subgroupMin(min_scored_point.score);
+        bool is_min_thread = min_scored_point.score == subgroup_smallest;
+        uvec4 subgroup_min_threads = subgroupBallot(is_min_thread);
+        uint subgroup_min_thread = subgroupBallotFindLSB(subgroup_min_threads);
+        min_index = subgroupShuffle(min_index, subgroup_min_thread);
+        
+        if (subgroupElect()) {
+            candidates[min_index] = scored_point;
+        }
+    } else {
+        if (subgroupElect()) {
+            candidates[candidates_count] = scored_point;
+        }
+        candidates_count++;
+    }
+    return true;
+}
+
+#endif

lib/segment/src/index/hnsw_index/gpu/shaders/iterators.comp

@@ -0,0 +1,27 @@
+#ifndef ITERATORS_HEADER
+#define ITERATORS_HEADER
+
+struct Iterator {
+    uint begin;
+    uint end;
+};
+
+#define NEW_ITERATOR(BEGIN, END) \
+    { \
+        (BEGIN), \
+        (END) \
+    }
+
+#define FOREACH_ITERATOR_VARIABLE_NAME _foreach_iterator
+
+#define FOREACH(ITERATOR) \
+    for ( \
+        Iterator FOREACH_ITERATOR_VARIABLE_NAME = ITERATOR; \
+        FOREACH_ITERATOR_VARIABLE_NAME.begin < FOREACH_ITERATOR_VARIABLE_NAME.end; \
+        FOREACH_ITERATOR_VARIABLE_NAME.begin += 1 \
+    )
+
+#define FOREACH_VALUE(CONTAINER) \
+    CONTAINER[FOREACH_ITERATOR_VARIABLE_NAME.begin]
+
+#endif

lib/segment/src/index/hnsw_index/gpu/shaders/links.comp

@@ -0,0 +1,46 @@
+#ifndef LINKS_HEADER
+#define LINKS_HEADER
+
+#ifndef LINKS_LAYOUT_SET
+#error "LINKS_LAYOUT_SET not defined"
+#endif
+
+#include "common.comp"
+#include "iterators.comp"
+
+layout(set = LINKS_LAYOUT_SET, binding = 0)
+readonly uniform LinksParams {
+    uint m;
+} links_params;
+
+layout(set = LINKS_LAYOUT_SET, binding = 1)
+buffer Links {
+    POINT_ID data[];
+} links;
+
+#define LINKS_CONTAINER links.data
+
+#define LEVEL_M links_params.m
+
+#define LINKS_CONTAINER_OFFSET(ID) \
+    ID * (LINKS_CAPACITY + 1)
+
+#define LINKS_NEW_ITERATOR(ID) \
+    NEW_ITERATOR((LINKS_CONTAINER_OFFSET(ID) + 1), (LINKS_CONTAINER_OFFSET(ID) + 1 + LINKS_COUNT(ID)))
+
+#define LINKS_COUNT(ID) \
+    LINKS_CONTAINER[LINKS_CONTAINER_OFFSET(ID)]
+
+#define GET_LINK(ID, I) \
+    LINKS_CONTAINER[LINKS_CONTAINER_OFFSET(ID) + 1 + I]
+
+#define LINKS_FOREACH(ID) \
+    FOREACH(LINKS_NEW_ITERATOR(ID))
+
+#define LINKS_FOREACH_VALUE \
+    FOREACH_VALUE(LINKS_CONTAINER)
+
+#define LINKS_SET_SIZE(ID, SIZE) \
+    LINKS_CONTAINER[LINKS_CONTAINER_OFFSET(ID)] = SIZE
+
+#endif

lib/segment/src/index/hnsw_index/gpu/shaders/nearest_heap.comp

@@ -0,0 +1,205 @@
+#ifndef NEAREST_HEAP_HEADER
+#define NEAREST_HEAP_HEADER
+
+#include "common.comp"
+
+shared ScoredPoint nearest[NEAREST_HEAP_CAPACITY];
+
+ScoredPoint nearest_smaller;
+uint nearest_smaller_index;
+uint nearest_count;
+
+void nearest_heap_init() {
+    nearest_count = 0;
+    nearest_smaller_index = 0;
+    nearest_smaller = ScoredPoint(UINT_MAX, positive_infinity);
+
+    uint capacity = NEAREST_HEAP_CAPACITY;
+    for (uint i = gl_SubgroupInvocationID; i < capacity; i += SUBGROUP_SIZE) {
+        uint buffer_index = i;
+        nearest[buffer_index] = ScoredPoint(UINT_MAX, negative_infinity);
+    }
+}
+
+// Push the nearest `ScoredPoint`. Returns true if the `ScoredPoint` was added.
+bool push_nearest(ScoredPoint scored_point) {
+    uint ef = NEAREST_HEAP_EF;
+    if (nearest_count == ef) {
+        if (scored_point.score > nearest_smaller.score) {
+            if (subgroupElect()) {
+                uint buffer_index = nearest_smaller_index;
+                nearest[buffer_index] = scored_point;
+            }
+            nearest_smaller_index = 0;
+            nearest_smaller.score = positive_infinity;
+            for (uint i = gl_SubgroupInvocationID; i < ef; i += SUBGROUP_SIZE) {
+                uint buffer_index = i;
+                float score = nearest[buffer_index].score;
+                if (score < nearest_smaller.score) {
+                    nearest_smaller.score = score;
+                    nearest_smaller_index = i;
+                }
+            }
+
+            float subgroup_smaller = subgroupMin(nearest_smaller.score);
+            bool is_min_thread = nearest_smaller.score == subgroup_smaller;
+            uvec4 subgroup_min_threads = subgroupBallot(is_min_thread);
+            uint subgroup_min_thread = subgroupBallotFindLSB(subgroup_min_threads);
+            nearest_smaller_index = subgroupShuffle(nearest_smaller_index, subgroup_min_thread);
+            nearest_smaller.score = subgroupShuffle(nearest_smaller.score, subgroup_min_thread);
+            nearest_smaller.id = nearest[nearest_smaller_index].id;
+
+            return true;
+        } else {
+            return false;
+        }
+    } else {
+        if (subgroupElect()) {
+            uint buffer_index = nearest_count;
+            nearest[buffer_index] = scored_point;
+        }
+        if (scored_point.score < nearest_smaller.score) {
+            nearest_smaller = scored_point;
+            nearest_smaller_index = nearest_count;
+        }
+        nearest_count++;
+        return true;
+    }
+}
+
+uint subgroup_bubble_sort(float score) {
+    uint invocation = gl_SubgroupInvocationID;
+    bool is_even = invocation % 2 == 0;
+    uint invocation_forward = is_even ? invocation + 1 : invocation - 1;
+    uint invocation_backward = invocation;
+    if (invocation != SUBGROUP_SIZE - 1 && invocation != 0) {
+        invocation_backward = is_even ? invocation - 1 : invocation + 1;
+    }
+
+    uint index = invocation;
+    for (uint j = 0; j < SUBGROUP_SIZE / 2; j++) {
+        bool is_even = gl_SubgroupInvocationID % 2 == 0;
+        bool was_changed = false;
+
+        // step one: compare with the next element
+        {
+            float next_score = subgroupShuffle(score, invocation_forward);
+            uint next_index = subgroupShuffle(index, invocation_forward);
+            bool need_exchange = is_even ? score < next_score : score > next_score;
+            if (need_exchange) {
+                was_changed = true;
+                score = next_score;
+                index = next_index;
+            }
+        }
+
+        // step two: compare with previous element
+        {
+            float next_score = subgroupShuffle(score, invocation_backward);
+            uint next_index = subgroupShuffle(index, invocation_backward);
+            bool need_exchange = is_even ? score > next_score : score < next_score;
+            if (need_exchange) {
+                was_changed = true;
+                score = next_score;
+                index = next_index;
+            }
+        }
+
+        if (!subgroupAny(was_changed)) {
+            break;
+        }
+    }
+    return index;
+}
+
+void sort_nearest_bubble() {
+    uint iterations_count = nearest_count / (SUBGROUP_SIZE / 2) + uint(sign(float(nearest_count % (SUBGROUP_SIZE / 2))));
+    for (uint _cnt = 0; _cnt < iterations_count; _cnt++) {
+        // do bubble sort in each subgroup to reduce the number of memory accesses
+        for (uint i = gl_SubgroupInvocationID; i < NEAREST_HEAP_CAPACITY; i += SUBGROUP_SIZE) {
+            ScoredPoint scored_point = nearest[i];
+
+            if (i >= nearest_count) {
+                scored_point.score = negative_infinity;
+            }
+            uint permut = subgroup_bubble_sort(scored_point.score);
+            float shuffled_scored_point_score = subgroupShuffle(scored_point.score, permut);
+            uint shuffled_scored_point_id = subgroupShuffle(scored_point.id, permut);
+            nearest[i] = ScoredPoint(shuffled_scored_point_id, shuffled_scored_point_score);
+            subgroupMemoryBarrierShared();
+        }
+
+        if (NEAREST_HEAP_CAPACITY > SUBGROUP_SIZE) {
+            // do bubble sort in each subgroup to reduce the number of memory accesses
+            for (uint i = gl_SubgroupInvocationID + (SUBGROUP_SIZE / 2); i < NEAREST_HEAP_CAPACITY - (SUBGROUP_SIZE / 2); i += SUBGROUP_SIZE) {
+                ScoredPoint scored_point = nearest[i];
+                
+                if (i >= nearest_count) {
+                    scored_point.score = negative_infinity;
+                }
+                uint permut = subgroup_bubble_sort(scored_point.score);
+                float shuffled_scored_point_score = subgroupShuffle(scored_point.score, permut);
+                uint shuffled_scored_point_id = subgroupShuffle(scored_point.id, permut);
+                nearest[i] = ScoredPoint(shuffled_scored_point_id, shuffled_scored_point_score);
+                subgroupMemoryBarrierShared();
+            }
+        }
+    }
+}
+
+void bitonic_cmp(uint i, uint j) {
+    uint b_i = i;
+    uint b_j = j;
+    if (i < j && j < nearest_count && nearest[b_j].score > nearest[b_i].score) {
+        ScoredPoint temp = nearest[b_i];
+        nearest[b_i] = nearest[b_j];
+        nearest[b_j] = temp;
+    }
+}
+
+void sort_nearest() {
+    for (uint group_offset = 0; group_offset < nearest_count; group_offset += SUBGROUP_SIZE) {
+        ScoredPoint scored_point = ScoredPoint(0, negative_infinity);
+        uint i = group_offset + gl_SubgroupInvocationID;
+        if (i < nearest_count) {
+            scored_point = nearest[i];
+        }
+
+        uint permut = subgroup_bubble_sort(scored_point.score);
+        scored_point = ScoredPoint(
+            subgroupShuffle(scored_point.id, permut),
+            subgroupShuffle(scored_point.score, permut)
+        );
+
+        if (i < nearest_count) {
+            nearest[i] = scored_point;
+        }
+    }
+
+    for (uint k = 64; (k >> 1) < nearest_count; k <<= 1) {
+        for (uint i = gl_SubgroupInvocationID; i < nearest_count; i += SUBGROUP_SIZE) {
+            bitonic_cmp(i, i ^ (k - 1));
+        }
+        subgroupMemoryBarrierShared();
+
+        for (uint j = k >> 1; 0 < j; j >>= 1) {
+            for (uint i = gl_SubgroupInvocationID; i < nearest_count; i += SUBGROUP_SIZE) {
+                bitonic_cmp(i, i ^ j);
+            }
+            subgroupMemoryBarrierShared();
+        }
+    }
+}
+
+ScoredPoint get_nearest(uint i) {
+    return nearest[i];
+}
+
+void set_nearest(uint i, ScoredPoint scored_point) {
+    if (subgroupElect()) {
+        nearest[i] = scored_point;
+    }
+    subgroupMemoryBarrierShared();
+}
+
+#endif

lib/segment/src/index/hnsw_index/gpu/shaders/run_get_patch.comp

@@ -0,0 +1,139 @@
+#version 450
+
+#include "extensions.comp"
+
+#include "common.comp"
+layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#define VECTOR_STORAGE_LAYOUT_SET 1
+#include "vector_storage.comp"
+
+#define LINKS_LAYOUT_SET 2
+#include "links.comp"
+
+#define VISITED_FLAGS_LAYOUT_SET 3
+#include "visited_flags.comp"
+
+#include "search_context.comp"
+
+struct Request {
+    uint id;
+    uint entry;
+};
+
+layout(set = 0, binding = 0) buffer Requests {
+    Request data[];
+} requests;
+
+layout(set = 0, binding = 1) buffer Patches {
+    uint data[];
+} patches;
+
+layout(set = 0, binding = 2) buffer NewEntries {
+    uint data[];
+} new_entries;
+
+void clear_patches() {
+    uint subgroup_index = SUBGROUP_ID;
+    uint patch_size = LEVEL_M + 2; // id + links_count + links
+    uint all_patches_size = (LEVEL_M + 1) * patch_size;
+    uint patches_offset = subgroup_index * all_patches_size;
+    for (uint i = gl_SubgroupInvocationID; i < all_patches_size; i += SUBGROUP_SIZE) {
+        patches.data[patches_offset + i] = UINT_MAX;
+    }
+    subgroupMemoryBarrierShared();
+}
+
+void main() {
+    clear_patches();
+    search_context_init();
+    uint subgroup_index = SUBGROUP_ID;
+    Request request = requests.data[subgroup_index];
+
+    set_target(request.id);
+    float entry_score = similarity(request.entry);
+    uint search_count = search(request.id, ScoredPoint(request.entry, entry_score));
+    if (search_count == UINT_MAX) {
+        // error while searching. propably candidates heap overflow
+        return;
+    }
+
+    if (subgroupElect()) {
+        if (search_count > 0) {
+            new_entries.data[subgroup_index] = get_nearest(0).id;
+        } else {
+            new_entries.data[subgroup_index] = request.entry;
+        }
+    }
+
+    uint count = run_heuristic();
+
+    uint patch_size = LEVEL_M + 2; // id + links_count + links
+    uint all_patches_size = (LEVEL_M + 1) * patch_size;
+    uint patches_offset = subgroup_index * all_patches_size;
+    
+    // write request.id patch
+    if (subgroupElect()) {
+        patches.data[patches_offset] = request.id;
+        patches.data[patches_offset + 1] = count;
+        for (uint i = 0; i < count; i++) {
+            patches.data[patches_offset + i + 2] = get_nearest(i).id;
+        }
+    }
+    subgroupMemoryBarrierShared();
+
+    for (uint i = 0; i < count; i++) {
+        uint other_id = patches.data[patches_offset + i + 2];
+        uint other_links_count = LINKS_COUNT(other_id);
+        uint other_patches_offset = patches_offset + (i + 1) * patch_size;
+        if (other_links_count < LEVEL_M) {
+            if (subgroupElect()) {
+                patches.data[other_patches_offset] = other_id;
+                other_patches_offset++;
+                patches.data[other_patches_offset] = other_links_count + 1;
+                other_patches_offset++;
+                LINKS_FOREACH(other_id) {
+                    patches.data[other_patches_offset] = LINKS_FOREACH_VALUE;
+                    other_patches_offset++;
+                }
+                patches.data[other_patches_offset] = request.id;
+            }
+        } else {
+            set_target(other_id);
+            nearest_count = 0;
+            LINKS_FOREACH(other_id) {
+                set_nearest(
+                    nearest_count,
+                    ScoredPoint(
+                        LINKS_FOREACH_VALUE,
+                        similarity(LINKS_FOREACH_VALUE))
+                );
+                nearest_count++;
+            }
+            set_nearest(
+                nearest_count,
+                ScoredPoint(
+                    request.id,
+                    similarity(request.id))
+            );
+            nearest_count++;
+
+            sort_nearest();
+            uint other_new_links_count = run_heuristic();
+
+            if (subgroupElect()) {
+                patches.data[other_patches_offset] = other_id;
+                other_patches_offset++;
+                patches.data[other_patches_offset] = other_new_links_count;
+                other_patches_offset++;
+            }
+            for (uint j = 0; j < other_new_links_count; j++) {
+                if (subgroupElect()) {
+                    uint l = get_nearest(j).id;
+                    patches.data[other_patches_offset] = l;
+                    other_patches_offset++;
+                }
+            }
+        }
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/run_greedy_search.comp

@@ -0,0 +1,44 @@
+#version 450
+
+#include "extensions.comp"
+
+#include "common.comp"
+layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#define VECTOR_STORAGE_LAYOUT_SET 1
+#include "vector_storage.comp"
+
+#define LINKS_LAYOUT_SET 2
+#include "links.comp"
+
+#define VISITED_FLAGS_LAYOUT_SET 3
+#include "visited_flags.comp"
+
+#include "search_context.comp"
+
+struct SearchRequest {
+    uint id;
+    uint entry;
+};
+
+layout(set = 0, binding = 0) buffer SearchRequests {
+    SearchRequest data[];
+} search_requests;
+
+layout(set = 0, binding = 1) buffer SearchResults {
+    uint data[];
+} search_results;
+
+void main() {
+    search_context_init();
+    uint subgroup_index = SUBGROUP_ID;
+    SearchRequest search_request = search_requests.data[subgroup_index];
+
+    set_target(search_request.id);
+    float entry_score = similarity(search_request.entry);
+    ScoredPoint result = greedy_search(search_request.id, ScoredPoint(search_request.entry, entry_score));
+
+    if (subgroupElect()) {
+        search_results.data[subgroup_index] = result.id;
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/run_insert_vector.comp

@@ -0,0 +1,129 @@
+#version 450
+
+#include "extensions.comp"
+
+#include "common.comp"
+layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#define VECTOR_STORAGE_LAYOUT_SET 1
+#include "vector_storage.comp"
+
+#define LINKS_LAYOUT_SET 2
+#include "links.comp"
+
+#define VISITED_FLAGS_LAYOUT_SET 3
+#include "visited_flags.comp"
+
+#include "search_context.comp"
+
+struct Request {
+    uint id;
+    uint entry;
+};
+
+layout(set = 0, binding = 0) buffer Requests {
+    Request data[];
+} requests;
+
+layout(set = 0, binding = 1) buffer NewEntries {
+    uint data[];
+} new_entries;
+
+layout(set = 0, binding = 2) buffer Atomics {
+    uint data[];
+} atomics;
+
+void main() {
+    search_context_init();
+    uint subgroup_index = SUBGROUP_ID;
+    Request request = requests.data[subgroup_index];
+
+    set_target(request.id);
+    float entry_score = similarity(request.entry);
+#ifdef NO_EQUIVALENCY
+    ScoredPoint new_entry = greedy_search(request.id, ScoredPoint(request.entry, entry_score));
+    uint search_count = search(request.id, new_entry);
+#else
+    uint search_count = search(request.id, ScoredPoint(request.entry, entry_score));
+#endif
+    if (search_count == UINT_MAX) {
+        // error while searching. propably candidates heap overflow
+        return;
+    }
+
+    if (subgroupElect()) {
+        if (search_count > 0) {
+            new_entries.data[subgroup_index] = get_nearest(0).id;
+        } else {
+            new_entries.data[subgroup_index] = request.entry;
+        }
+    }
+
+    uint count = run_heuristic();
+
+    // sync all threads to make sure that we don't do search while updating links
+    barrier();
+
+    // update request.id links
+    for (uint i = gl_SubgroupInvocationID; i < count; i += SUBGROUP_SIZE) {
+        GET_LINK(request.id, i) = get_nearest(i).id;
+    }
+    if (subgroupElect()) {
+        LINKS_SET_SIZE(request.id, count);
+    }
+
+    subgroupMemoryBarrier();
+
+    uint generation = VISITED_GENERATION;
+
+    for (uint i = 0; i < count; i++) {
+        uint other_id = GET_LINK(request.id, i);
+
+        bool is_changed = false;
+        if (subgroupElect()) {
+            uint other_atomic = atomicExchange(atomics.data[other_id], generation);
+            is_changed = other_atomic == generation;
+        }
+        if (subgroupAny(is_changed)) {
+            // other_id is already being processed by another subgroup
+            continue;
+        }
+
+        uint other_links_count = LINKS_COUNT(other_id);
+        if (other_links_count < LEVEL_M) {
+            if (subgroupElect()) {
+                GET_LINK(other_id, other_links_count) = request.id;
+                LINKS_SET_SIZE(other_id, other_links_count + 1);
+            }
+        } else {
+            set_target(other_id);
+            nearest_count = 0;
+            LINKS_FOREACH(other_id) {
+                set_nearest(
+                    nearest_count,
+                    ScoredPoint(
+                        LINKS_FOREACH_VALUE,
+                        similarity(LINKS_FOREACH_VALUE))
+                );
+                nearest_count++;
+            }
+            set_nearest(
+                nearest_count,
+                ScoredPoint(
+                    request.id,
+                    similarity(request.id))
+            );
+            nearest_count++;
+
+            sort_nearest();
+            uint other_new_links_count = run_heuristic();
+
+            for (uint j = gl_SubgroupInvocationID; j < other_new_links_count; j += SUBGROUP_SIZE) {
+                GET_LINK(other_id, j) = get_nearest(j).id;
+            }
+            if (subgroupElect()) {
+                LINKS_SET_SIZE(other_id, other_new_links_count);
+            }
+        }
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/search_context.comp

@@ -0,0 +1,140 @@
+#ifndef SEARCH_CONTEXT_HEADER
+#define SEARCH_CONTEXT_HEADER
+
+#include "common.comp"
+#include "vector_storage.comp"
+#include "links.comp"
+#include "nearest_heap.comp"
+#include "candidates_heap.comp"
+#include "visited_flags.comp"
+
+void process_candidate(ScoredPoint scored_point) {
+    bool not_presented = nearest_smaller.id != scored_point.id;
+    bool was_added = push_nearest(scored_point);
+    if (was_added && not_presented) {
+        push_candidate(scored_point);
+        subgroupMemoryBarrierShared();
+    }
+}
+
+void process_candidate(POINT_ID candidate_id) {
+    uint links_count = LINKS_COUNT(candidate_id);
+    uint subgroup_link = 0;
+    bool subgroup_is_visited = true;
+    uint invocation = gl_SubgroupInvocationID;
+
+    for (uint i = 0; i < links_count; i++) {
+        if (i % SUBGROUP_SIZE == 0 && i + invocation < links_count) {
+            subgroup_link = GET_LINK(candidate_id, i + invocation);
+            subgroup_is_visited = is_visited(subgroup_link);
+            if (!subgroup_is_visited) {
+                mark_visited(subgroup_link);
+            }
+        }
+
+        uint link = subgroupShuffle(subgroup_link, i % SUBGROUP_SIZE);
+        bool is_visited = subgroupShuffle(subgroup_is_visited, i % SUBGROUP_SIZE);
+
+        if (!is_visited) {
+            float score = similarity(link);
+            ScoredPoint scored_point = ScoredPoint(link, score);
+            process_candidate(scored_point);
+        }
+    }
+
+    // flush all changed visited flags
+    subgroupMemoryBarrierBuffer();
+}
+
+ScoredPoint greedy_search(POINT_ID point_id, ScoredPoint entry) {
+    bool changed = true;
+    ScoredPoint result = entry;
+
+    uint subgroup_link = 0;
+    uint invocation = gl_SubgroupInvocationID;
+
+    while (changed) {
+        changed = false;
+
+        uint current = result.id;
+        uint links_count = LINKS_COUNT(current);
+        for (uint i = 0; i < links_count; i++) {
+            if (i % SUBGROUP_SIZE == 0 && i + invocation < links_count) {
+                subgroup_link = GET_LINK(current, i + invocation);
+            }
+            uint link = subgroupShuffle(subgroup_link, i % SUBGROUP_SIZE);
+
+            float score = similarity(link);
+            if (score > result.score) {
+                changed = true;
+                result = ScoredPoint(link, score);
+            }
+        }
+    }
+    return result;
+}
+
+uint search(POINT_ID point_id, ScoredPoint entry) {
+#ifndef EXACT
+    entry = greedy_search(point_id, entry);
+#endif
+
+    if (subgroupElect()) {
+        mark_visited(point_id);
+        mark_visited(entry.id);
+    }
+    push_nearest(entry);
+    push_candidate(entry);
+    groupMemoryBarrier();
+    
+    while (candidates_count > 0) {
+        ScoredPoint candidate = pop_candidate();
+        if (candidate.score < nearest_smaller.score) {
+            continue;
+        }
+        process_candidate(candidate.id);
+    }
+
+    process_candidate(point_id);
+
+    sort_nearest();
+    return nearest_count;
+}
+
+uint run_heuristic() {
+    uint result_count = 0;
+
+    for (uint i = 0; i < nearest_count; i++) {
+        if (result_count >= LEVEL_M) {
+            break;
+        }
+
+        ScoredPoint current_closest = get_nearest(i);
+        set_target(current_closest.id);
+        bool is_good = true;
+
+        for (uint j = 0; j < result_count; j++) {
+            ScoredPoint selected_point = get_nearest(j);
+            float dist_to_already_selected = similarity(selected_point.id);
+            if (dist_to_already_selected > current_closest.score) {
+                is_good = false;
+                break;
+            }
+        }
+
+        if (is_good) {
+            set_nearest(result_count, current_closest);
+            result_count += 1;
+        }
+    }
+
+    nearest_count = result_count;
+    return result_count;
+}
+
+void search_context_init() {
+    nearest_heap_init();
+    candidates_heap_init();
+}
+
+#endif

lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_candidates_heap.comp

@@ -0,0 +1,51 @@
+#version 450
+
+#include "extensions.comp"
+
+#include "common.comp"
+layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#include "candidates_heap.comp"
+
+layout(set = 0, binding = 0)
+readonly uniform TestParams {
+    uint input_counts;
+} test_params;
+
+layout(set = 0, binding = 1) buffer ScoredPointsInput {
+    ScoredPoint data[];
+} input_points;
+
+layout(set = 0, binding = 2) buffer ScoresOuput {
+    ScoredPoint data[];
+} scores_output;
+
+void main() {
+    candidates_heap_init();
+    uint subgroup_index = SUBGROUP_ID;
+
+    for (uint i = 0; i < test_params.input_counts; i++) {
+        uint point_index = subgroup_index * test_params.input_counts + i;
+        ScoredPoint point = input_points.data[point_index];
+        push_candidate(point);
+    }
+
+    while (candidates_count > 0) {
+        ScoredPoint candidate = pop_candidate();
+    }
+
+    for (uint i = 0; i < test_params.input_counts; i++) {
+        uint point_index = subgroup_index * test_params.input_counts + i;
+        ScoredPoint point = input_points.data[point_index];
+        push_candidate(point);
+    }
+
+    uint output_index = subgroup_index * test_params.input_counts;
+    while (candidates_count > 0) {
+        ScoredPoint candidate = pop_candidate();
+        if (subgroupElect()) {
+            scores_output.data[output_index] = candidate;
+        }
+        output_index++;
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_heuristic.comp

@@ -0,0 +1,50 @@
+#version 450
+
+#include "extensions.comp"
+
+#include "common.comp"
+layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#define VECTOR_STORAGE_LAYOUT_SET 1
+#include "vector_storage.comp"
+
+#define LINKS_LAYOUT_SET 2
+#include "links.comp"
+
+#define VISITED_FLAGS_LAYOUT_SET 3
+#include "visited_flags.comp"
+
+#include "search_context.comp"
+
+struct SearchRequest {
+    uint id;
+    uint entry;
+};
+
+layout(set = 0, binding = 0) buffer SearchRequests {
+    SearchRequest data[];
+} search_requests;
+
+layout(set = 0, binding = 1) buffer SearchResults {
+    ScoredPoint data[];
+} search_results;
+
+void main() {
+    search_context_init();
+    uint subgroup_index = SUBGROUP_ID;
+    SearchRequest search_request = search_requests.data[subgroup_index];
+
+    set_target(search_request.id);
+    float entry_score = similarity(search_request.entry);
+    search(search_request.id, ScoredPoint(search_request.entry, entry_score));
+    uint count = run_heuristic();
+
+    barrier();
+
+    if (subgroupElect()) {
+        uint result_index_offset = NEAREST_HEAP_EF * subgroup_index;
+        for (uint i = 0; i < count; i++) {
+            search_results.data[result_index_offset + i] = get_nearest(i);
+        }
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_hnsw_search.comp

@@ -0,0 +1,53 @@
+#version 450
+
+#include "extensions.comp"
+
+#include "common.comp"
+layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#define VECTOR_STORAGE_LAYOUT_SET 1
+#include "vector_storage.comp"
+
+#define LINKS_LAYOUT_SET 2
+#include "links.comp"
+
+#define VISITED_FLAGS_LAYOUT_SET 3
+#include "visited_flags.comp"
+
+#include "search_context.comp"
+
+struct SearchRequest {
+    uint id;
+    uint entry;
+};
+
+layout(set = 0, binding = 0) buffer SearchRequests {
+    SearchRequest data[];
+} search_requests;
+
+layout(set = 0, binding = 1) buffer SearchResults {
+    ScoredPoint data[];
+} search_results;
+
+void main() {
+    search_context_init();
+    uint subgroup_index = SUBGROUP_ID;
+    SearchRequest search_request = search_requests.data[subgroup_index];
+
+    set_target(search_request.id);
+    float entry_score = similarity(search_request.entry);
+    uint count = search(search_request.id, ScoredPoint(search_request.entry, entry_score));
+
+    barrier();
+
+    if (subgroupElect()) {
+        uint result_index_offset = NEAREST_HEAP_EF * subgroup_index;
+        uint i = 0;
+        for (; i < count; i++) {
+            search_results.data[result_index_offset + i] = get_nearest(i);
+        }
+        for (; i < NEAREST_HEAP_EF; i++) {
+            search_results.data[result_index_offset + i] = ScoredPoint(UINT_MAX, 0.0);
+        }
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/tests/test_nearest_heap.comp

@@ -0,0 +1,58 @@
+#version 450
+
+#include "extensions.comp"
+
+#include "common.comp"
+layout(local_size_x = SUBGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#include "nearest_heap.comp"
+
+layout(set = 0, binding = 0)
+readonly uniform TestParams {
+    uint input_counts;
+} test_params;
+
+layout(set = 0, binding = 1) buffer ScoredPointsInput {
+    ScoredPoint data[];
+} input_points;
+
+layout(set = 0, binding = 2) buffer ScoresOuput {
+    float data[];
+} scores_output;
+
+layout(set = 0, binding = 3) buffer SortedOuput {
+    uint data[];
+} sorted_output;
+
+void main() {
+    nearest_heap_init();
+    uint subgroup_index = SUBGROUP_ID;
+
+    for (uint i = 0; i < test_params.input_counts; i++) {
+        uint point_index = subgroup_index * test_params.input_counts + i;
+        ScoredPoint point = input_points.data[point_index];
+        push_nearest(point);
+
+        groupMemoryBarrier();
+        memoryBarrier();
+
+        float nearest_smaller_score_check = subgroupMin(nearest_smaller.score);
+        bool is_different = nearest_smaller_score_check != nearest_smaller.score;
+        if (subgroupAny(is_different)) {
+            if (subgroupElect()) {
+                scores_output.data[point_index] = positive_infinity;
+            }
+        } else if (subgroupElect()) {
+            scores_output.data[point_index] = nearest_smaller.score;
+        }
+    }
+
+    sort_nearest();
+    if (subgroupElect()) {
+        for (uint i = 0; i < NEAREST_HEAP_EF; i++) {
+            uint point_index = subgroup_index * NEAREST_HEAP_EF + i;
+            ScoredPoint point = get_nearest(i);
+            sorted_output.data[point_index] = point.id;
+        }
+    }
+}

lib/segment/src/index/hnsw_index/gpu/shaders/visited_flags.comp

@@ -0,0 +1,37 @@
+#ifndef VISITED_FLAGS_HEADER
+#define VISITED_FLAGS_HEADER
+
+#ifndef VISITED_FLAGS_LAYOUT_SET
+#error "VISITED_FLAGS_LAYOUT_SET not defined"
+#endif
+
+#include "common.comp"
+#include "vector_storage.comp"
+
+layout(set = VISITED_FLAGS_LAYOUT_SET, binding = 0)
+readonly uniform VisitedFlagsParams {
+    uint generation;
+} visited_flags_params;
+
+layout(set = VISITED_FLAGS_LAYOUT_SET, binding = 1)
+buffer VisitedFlagsBuffer {
+    uint8_t data[];
+} visited_flags;
+
+#define BYTE_MASK 0xFF
+#define VISITED_GENERATION visited_flags_params.generation
+
+bool is_visited(POINT_ID point_id) {
+    uint subgroup_index = SUBGROUP_ID;
+    uint index = subgroup_index * VISITED_FLAGS_CAPACITY + point_id % VISITED_FLAGS_CAPACITY;
+    uint prev_generation = uint(visited_flags.data[index]);
+    return prev_generation == VISITED_GENERATION;
+}
+
+void mark_visited(POINT_ID point_id) {
+    uint subgroup_index = SUBGROUP_ID;
+    uint index = subgroup_index * VISITED_FLAGS_CAPACITY + point_id % VISITED_FLAGS_CAPACITY;
+    visited_flags.data[index] = uint8_t(VISITED_GENERATION);
+}
+
+#endif

lib/segment/src/index/hnsw_index/graph_layers.rs

@@ -137,6 +137,39 @@ pub trait GraphLayersBase {
         }
         current_point
     }
+
+    fn search_entry_on_level(
+        &self,
+        entry_point: PointOffsetType,
+        level: usize,
+        points_scorer: &mut FilteredScorer,
+    ) -> ScoredPointOffset {
+        let limit = self.get_m(level);
+        let mut links: Vec<PointOffsetType> = Vec::with_capacity(2 * self.get_m(0));
+        let mut current_point = ScoredPointOffset {
+            idx: entry_point,
+            score: points_scorer.score_point(entry_point),
+        };
+
+        let mut changed = true;
+        while changed {
+            changed = false;
+
+            links.clear();
+            self.links_map(current_point.idx, level, |link| {
+                links.push(link);
+            });
+
+            let scores = points_scorer.score_points(&mut links, limit);
+            scores.iter().copied().for_each(|score_point| {
+                if score_point.score > current_point.score {
+                    changed = true;
+                    current_point = score_point;
+                }
+            });
+        }
+        current_point
+    }
 }
 
 impl<TGraphLinks: GraphLinks> GraphLayersBase for GraphLayers<TGraphLinks> {

lib/segment/src/index/hnsw_index/graph_layers_builder.rs

@@ -6,6 +6,7 @@ use std::sync::atomic::AtomicUsize;
 use bitvec::prelude::BitVec;
 use common::fixed_length_priority_queue::FixedLengthPriorityQueue;
 use common::types::{PointOffsetType, ScoreType, ScoredPointOffset};
+use itertools::Itertools;
 use parking_lot::{Mutex, MutexGuard, RwLock};
 use rand::distributions::Uniform;
 use rand::Rng;
@@ -25,22 +26,22 @@ pub type LockedLayersContainer = Vec<LockedLinkContainer>;
 /// Same as `GraphLayers`,  but allows to build in parallel
 /// Convertible to `GraphLayers`
 pub struct GraphLayersBuilder {
-    max_level: AtomicUsize,
-    m: usize,
-    m0: usize,
-    ef_construct: usize,
+    pub max_level: AtomicUsize,
+    pub m: usize,
+    pub m0: usize,
+    pub ef_construct: usize,
     // Factor of level probability
-    level_factor: f64,
+    pub level_factor: f64,
     // Exclude points according to "not closer than base" heuristic?
-    use_heuristic: bool,
-    links_layers: Vec<LockedLayersContainer>,
-    entry_points: Mutex<EntryPoints>,
+    pub use_heuristic: bool,
+    pub links_layers: Vec<LockedLayersContainer>,
+    pub entry_points: Mutex<EntryPoints>,
 
     // Fields used on construction phase only
-    visited_pool: VisitedPool,
+    pub visited_pool: VisitedPool,
 
     // List of bool flags, which defines if the point is already indexed or not
-    ready_list: RwLock<BitVec>,
+    pub ready_list: RwLock<BitVec>,
 }
 
 impl GraphLayersBase for GraphLayersBuilder {
@@ -206,7 +207,7 @@ impl GraphLayersBuilder {
         picked_level.round() as usize
     }
 
-    fn get_point_level(&self, point_id: PointOffsetType) -> usize {
+    pub fn get_point_level(&self, point_id: PointOffsetType) -> usize {
         self.links_layers[point_id as usize].len() - 1
     }
 
@@ -286,7 +287,7 @@ impl GraphLayersBuilder {
     }
 
     /// <https://github.com/nmslib/hnswlib/issues/99>
-    fn select_candidates_with_heuristic<F>(
+    pub(crate) fn select_candidates_with_heuristic<F>(
         candidates: FixedLengthPriorityQueue<ScoredPointOffset>,
         m: usize,
         score_internal: F,
@@ -321,12 +322,14 @@ impl GraphLayersBuilder {
                     // Let's find closest one on same level
 
                     // greedy search for a single closest point
-                    self.search_entry(
+                    let new_entry = self.search_entry(
                         entry_point.point_id,
                         entry_point.level,
                         level,
                         &mut points_scorer,
-                    )
+                    );
+
+                    new_entry
                 } else {
                     ScoredPointOffset {
                         idx: entry_point.point_id,
@@ -477,6 +480,128 @@ impl GraphLayersBuilder {
             sum as f32 / count as f32
         }
     }
+
+    pub fn get_patch(
+        &self,
+        request: super::gpu::gpu_search_context::GpuRequest,
+        level: usize,
+        points_scorer: &mut FilteredScorer,
+    ) -> (
+        Vec<super::gpu::gpu_search_context::GpuGraphLinksPatch>,
+        Vec<PointOffsetType>,
+    ) {
+        let entry = request.entry;
+        let point_id = request.id;
+        let mut patches = vec![];
+
+        let level_m = self.get_m(level);
+        let mut visited_list = self.get_visited_list_from_pool();
+
+        visited_list.check_and_update_visited(entry);
+
+        let level_entry = ScoredPointOffset {
+            idx: entry,
+            score: points_scorer.score_point(entry),
+        };
+        let mut search_context = SearchContext::new(level_entry, self.ef_construct);
+
+        self._search_on_level(&mut search_context, level, &mut visited_list, points_scorer);
+
+        let new_entries = search_context
+            .nearest
+            .clone()
+            .into_vec()
+            .iter()
+            .map(|p| p.idx)
+            .collect_vec();
+
+        let scorer = |a, b| points_scorer.score_internal(a, b);
+
+        let selected_nearest = {
+            let selected_nearest =
+                Self::select_candidates_with_heuristic(search_context.nearest, level_m, scorer);
+            patches.push(super::gpu::gpu_search_context::GpuGraphLinksPatch {
+                id: point_id,
+                links: selected_nearest.clone(),
+            });
+            selected_nearest
+        };
+
+        for &other_point in &selected_nearest {
+            let mut other_point_links = self.links_layers[other_point as usize][level]
+                .read()
+                .clone();
+            if other_point_links.len() < level_m {
+                // If linked point is lack of neighbours
+                other_point_links.push(point_id);
+            } else {
+                let mut candidates = BinaryHeap::with_capacity(level_m + 1);
+                candidates.push(ScoredPointOffset {
+                    idx: point_id,
+                    score: scorer(point_id, other_point),
+                });
+                for other_point_link in other_point_links.iter().take(level_m).copied() {
+                    candidates.push(ScoredPointOffset {
+                        idx: other_point_link,
+                        score: scorer(other_point_link, other_point),
+                    });
+                }
+                let selected_candidates = Self::select_candidate_with_heuristic_from_sorted(
+                    candidates.into_sorted_vec().into_iter().rev(),
+                    level_m,
+                    scorer,
+                );
+                other_point_links.clear(); // this do not free memory, which is good
+                for selected in selected_candidates.iter().copied() {
+                    other_point_links.push(selected);
+                }
+            }
+            patches.push(super::gpu::gpu_search_context::GpuGraphLinksPatch {
+                id: other_point,
+                links: other_point_links,
+            });
+        }
+
+        (patches, new_entries)
+    }
+
+    pub fn try_apply_patch(
+        &self,
+        level: usize,
+        patches: Vec<super::gpu::gpu_search_context::GpuGraphLinksPatch>,
+    ) -> bool {
+        let mut locks = Vec::with_capacity(patches.len());
+        let mut links_to_apply = Vec::with_capacity(patches.len());
+        for patch in patches {
+            if let Some(lock) = self.links_layers[patch.id as usize][level].try_write() {
+                //if lock.as_slice() != patch.old_links.as_slice() {
+                //    return false;
+                //}
+
+                locks.push(lock);
+                links_to_apply.push(patch.links);
+            } else {
+                return false;
+            }
+        }
+
+        for (links, lock) in links_to_apply.into_iter().zip(locks.iter_mut()) {
+            **lock = links;
+        }
+
+        true
+    }
+
+    pub fn apply_patch(
+        &self,
+        level: usize,
+        patches: Vec<super::gpu::gpu_search_context::GpuGraphLinksPatch>,
+    ) {
+        for patch in patches {
+            let mut lock = self.links_layers[patch.id as usize][level].write();
+            *lock = patch.links;
+        }
+    }
 }
 
 #[cfg(test)]

lib/segment/src/index/hnsw_index/hnsw.rs

@@ -20,6 +20,7 @@ use rand::thread_rng;
 use rayon::prelude::*;
 use rayon::ThreadPool;
 
+use super::gpu::get_gpu_min_points_count;
 use super::graph_links::{GraphLinks, GraphLinksMmap};
 use crate::common::operation_error::{check_process_stopped, OperationError, OperationResult};
 use crate::common::operation_time_statistics::{
@@ -31,6 +32,10 @@ use crate::data_types::vectors::{QueryVector, VectorInternal, VectorRef};
 use crate::id_tracker::IdTrackerSS;
 use crate::index::hnsw_index::build_condition_checker::BuildConditionChecker;
 use crate::index::hnsw_index::config::HnswGraphConfig;
+use crate::index::hnsw_index::gpu::gpu_graph_builder::build_hnsw_on_gpu;
+use crate::index::hnsw_index::gpu::{
+    get_gpu_force_half_precision, get_gpu_indexing, get_gpu_max_groups,
+};
 use crate::index::hnsw_index::graph_layers::GraphLayers;
 use crate::index::hnsw_index::graph_layers_builder::GraphLayersBuilder;
 use crate::index::hnsw_index::point_scorer::FilteredScorer;
@@ -39,6 +44,7 @@ use crate::index::sample_estimation::sample_check_cardinality;
 use crate::index::struct_payload_index::StructPayloadIndex;
 use crate::index::visited_pool::{VisitedListHandle, VisitedPool};
 use crate::index::{PayloadIndex, VectorIndex};
+use crate::payload_storage::FilterContext;
 use crate::telemetry::VectorIndexSearchesTelemetry;
 use crate::types::Condition::Field;
 use crate::types::{
@@ -91,6 +97,7 @@ pub struct HnswIndexOpenArgs<'a> {
     pub payload_index: Arc<AtomicRefCell<StructPayloadIndex>>,
     pub hnsw_config: HnswConfig,
     pub permit: Option<Arc<CpuPermit>>,
+    pub gpu_device: Option<Arc<gpu::Device>>,
     pub stopped: &'a AtomicBool,
 }
 
@@ -104,6 +111,7 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
             payload_index,
             hnsw_config,
             permit,
+            gpu_device,
             stopped,
         } = args;
 
@@ -162,6 +170,7 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
                 &payload_index.borrow(),
                 &hnsw_config,
                 num_cpus,
+                gpu_device,
                 stopped,
             )?;
 
@@ -209,6 +218,7 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
         payload_index: &StructPayloadIndex,
         hnsw_config: &HnswConfig,
         num_cpus: usize,
+        gpu_device: Option<Arc<gpu::Device>>,
         stopped: &AtomicBool,
     ) -> OperationResult<(HnswGraphConfig, GraphLayers<TGraphLinks>)> {
         let total_vector_count = vector_storage.total_vector_count();
@@ -237,20 +247,22 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
         let mut rng = thread_rng();
         let deleted_bitslice = vector_storage.deleted_vector_bitslice();
 
+        // TODO(gpu): log device name
         debug!("building HNSW for {total_vector_count} vectors with {num_cpus} CPUs");
 
+        let num_entries = std::cmp::max(
+            1,
+            total_vector_count
+                .checked_div(full_scan_threshold)
+                .unwrap_or(0)
+                * 10,
+        );
         let mut graph_layers_builder = GraphLayersBuilder::new(
             total_vector_count,
             config.m,
             config.m0,
             config.ef_construct,
-            std::cmp::max(
-                1,
-                total_vector_count
-                    .checked_div(full_scan_threshold)
-                    .unwrap_or(0)
-                    * 10,
-            ),
+            num_entries,
             HNSW_USE_HEURISTIC,
         );
 
@@ -289,49 +301,108 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
         let mut indexed_vectors = 0;
 
         if config.m > 0 {
-            let mut ids_iterator = id_tracker.iter_ids_excluding(deleted_bitslice);
+            let timer = std::time::Instant::now();
 
-            let first_few_ids: Vec<_> = ids_iterator
-                .by_ref()
-                .take(SINGLE_THREADED_HNSW_BUILD_THRESHOLD)
-                .collect();
-            let ids: Vec<_> = ids_iterator.collect();
+            let use_gpu = total_vector_count > get_gpu_min_points_count() && get_gpu_indexing();
+            let gpu_constructed_graph = if let Some(gpu_device) = &gpu_device {
+                if use_gpu {
+                    let ids: Vec<_> = id_tracker.iter_ids_excluding(deleted_bitslice).collect();
+                    indexed_vectors = ids.len();
 
-            indexed_vectors = ids.len() + first_few_ids.len();
+                    let points_scorer_builder = |vector_id| {
+                        let vector = vector_storage.get_vector(vector_id);
+                        let vector = vector.as_vec_ref().into();
+                        let raw_scorer = if let Some(quantized_storage) = quantized_vectors.as_ref()
+                        {
+                            quantized_storage.raw_scorer(
+                                vector,
+                                id_tracker.deleted_point_bitslice(),
+                                vector_storage.deleted_vector_bitslice(),
+                                stopped,
+                            )
+                        } else {
+                            new_raw_scorer(
+                                vector,
+                                &vector_storage,
+                                id_tracker.deleted_point_bitslice(),
+                            )
+                        }?;
+                        Ok((raw_scorer, None))
+                    };
 
-            let insert_point = |vector_id| {
-                check_process_stopped(stopped)?;
-                let vector = vector_storage.get_vector(vector_id);
-                let vector = vector.as_vec_ref().into();
-                let raw_scorer = if let Some(quantized_storage) = quantized_vectors.as_ref() {
-                    quantized_storage.raw_scorer(
-                        vector,
-                        id_tracker.deleted_point_bitslice(),
-                        vector_storage.deleted_vector_bitslice(),
+                    Some(build_hnsw_on_gpu(
+                        gpu_device.clone(),
+                        &pool,
+                        &graph_layers_builder,
+                        get_gpu_max_groups(),
+                        &vector_storage,
+                        quantized_vectors.as_ref(),
+                        num_entries,
+                        get_gpu_force_half_precision(),
+                        SINGLE_THREADED_HNSW_BUILD_THRESHOLD,
+                        false,
+                        ids,
+                        points_scorer_builder,
                         stopped,
-                    )
+                    ))
                 } else {
-                    new_raw_scorer(vector, vector_storage, id_tracker.deleted_point_bitslice())
-                }?;
-                let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), None);
-
-                graph_layers_builder.link_new_point(vector_id, points_scorer);
-
-                // Ignore hardware counter, for internal operations
-                raw_scorer.take_hardware_counter().discard_results();
-
-                Ok::<_, OperationError>(())
+                    None
+                }
+            } else {
+                None
             };
 
-            for vector_id in first_few_ids {
-                insert_point(vector_id)?;
+            if let Some(Err(gpu_error)) = &gpu_constructed_graph {
+                log::warn!("Failed to build HNSW on GPU: {gpu_error}. Falling back to CPU.");
             }
 
-            if !ids.is_empty() {
-                pool.install(|| ids.into_par_iter().try_for_each(insert_point))?;
+            if let Some(Ok(gpu_graph)) = gpu_constructed_graph {
+                graph_layers_builder = gpu_graph;
+            } else {
+                let mut ids_iterator = id_tracker.iter_ids_excluding(deleted_bitslice);
+
+                let first_few_ids: Vec<_> = ids_iterator
+                    .by_ref()
+                    .take(SINGLE_THREADED_HNSW_BUILD_THRESHOLD)
+                    .collect();
+                let ids: Vec<_> = ids_iterator.collect();
+
+                indexed_vectors = ids.len() + first_few_ids.len();
+
+                let insert_point = |vector_id| {
+                    check_process_stopped(stopped)?;
+                    let vector = vector_storage.get_vector(vector_id);
+                    let vector = vector.as_vec_ref().into();
+                    let raw_scorer = if let Some(quantized_storage) = quantized_vectors.as_ref() {
+                        quantized_storage.raw_scorer(
+                            vector,
+                            id_tracker.deleted_point_bitslice(),
+                            vector_storage.deleted_vector_bitslice(),
+                            stopped,
+                        )
+                    } else {
+                        new_raw_scorer(vector, &vector_storage, id_tracker.deleted_point_bitslice())
+                    }?;
+                    let points_scorer = FilteredScorer::new(raw_scorer.as_ref(), None);
+
+                    graph_layers_builder.link_new_point(vector_id, points_scorer);
+
+                    // Ignore hardware counter, for internal operations
+                    raw_scorer.take_hardware_counter().discard_results();
+
+                    Ok::<_, OperationError>(())
+                };
+
+                for vector_id in first_few_ids {
+                    insert_point(vector_id)?;
+                }
+
+                if !ids.is_empty() {
+                    pool.install(|| ids.into_par_iter().try_for_each(insert_point))?;
+                }
             }
 
-            debug!("finish main graph");
+            debug!("finish main graph in time {:?}", timer.elapsed());
         } else {
             debug!("skip building main HNSW graph");
         }
@@ -390,6 +461,7 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
                         quantized_vectors,
                         payload_index,
                         &pool,
+                        gpu_device.clone(),
                         stopped,
                         &mut additional_graph,
                         payload_block.condition,
@@ -437,6 +509,7 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
         quantized_vectors: &Option<QuantizedVectors>,
         payload_index: &StructPayloadIndex,
         pool: &ThreadPool,
+        gpu_device: Option<Arc<gpu::Device>>,
         stopped: &AtomicBool,
         graph_layers_builder: &mut GraphLayersBuilder,
         condition: FieldCondition,
@@ -468,6 +541,63 @@ impl<TGraphLinks: GraphLinks> HNSWIndex<TGraphLinks> {
             }
         }
 
+        let use_gpu = points_to_index.len() > get_gpu_min_points_count() && get_gpu_indexing();
+        let gpu_constructed_graph = if let Some(gpu_device) = gpu_device {
+            if use_gpu {
+                let points_scorer_builder = |block_point_id| -> OperationResult<_> {
+                    let vector = vector_storage.get_vector(block_point_id);
+                    let vector = vector.as_vec_ref().into();
+                    let raw_scorer = match quantized_vectors.as_ref() {
+                        Some(quantized_storage) => quantized_storage.raw_scorer(
+                            vector,
+                            id_tracker.deleted_point_bitslice(),
+                            deleted_bitslice,
+                            stopped,
+                        ),
+                        None => new_raw_scorer(
+                            vector,
+                            vector_storage,
+                            id_tracker.deleted_point_bitslice(),
+                        ),
+                    }?;
+                    let block_condition_checker: Box<dyn FilterContext> =
+                        Box::new(BuildConditionChecker {
+                            filter_list: block_filter_list,
+                            current_point: block_point_id,
+                        });
+                    Ok((raw_scorer, Some(block_condition_checker)))
+                };
+                Some(build_hnsw_on_gpu(
+                    gpu_device.clone(),
+                    &pool,
+                    &graph_layers_builder,
+                    get_gpu_max_groups(),
+                    &vector_storage,
+                    quantized_vectors.as_ref(),
+                    1,
+                    get_gpu_force_half_precision(),
+                    SINGLE_THREADED_HNSW_BUILD_THRESHOLD,
+                    false,
+                    points_to_index.clone(),
+                    points_scorer_builder,
+                    stopped,
+                ))
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+
+        if let Some(Err(gpu_error)) = &gpu_constructed_graph {
+            log::warn!("Failed to build HNSW on GPU: {gpu_error}. Falling back to CPU.");
+        }
+
+        if let Some(Ok(gpu_graph)) = gpu_constructed_graph {
+            *graph_layers_builder = gpu_graph;
+            return Ok(());
+        }
+
         let insert_points = |block_point_id| {
             check_process_stopped(stopped)?;
 

lib/segment/src/index/hnsw_index/tests/test_graph_connectivity.rs

@@ -86,6 +86,7 @@ fn test_graph_connectivity() {
         payload_index: payload_index_ptr,
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/src/segment_constructor/segment_builder.rs

@@ -427,6 +427,23 @@ impl SegmentBuilder {
             id_tracker.versions_flusher()()?;
             let id_tracker_arc = Arc::new(AtomicRefCell::new(id_tracker));
 
+            let gpu_device = if crate::index::hnsw_index::gpu::get_gpu_indexing() {
+                crate::index::hnsw_index::gpu::GPU_DEVICES_MANAGER
+                    .as_ref()
+                    .map(|devices_manager| devices_manager.lock_device())
+                    .ok()
+                    .flatten()
+            } else {
+                None
+            };
+            /*
+            if let Some(_gpu_device) = &gpu_device {
+                if permit.num_cpus > 1 {
+                    permit.release_count(permit.num_cpus - 1);
+                }
+            }
+             */
+
             // Arc permit to share it with each vector store
             let permit = Arc::new(permit);
 
@@ -500,6 +517,9 @@ impl SegmentBuilder {
                     payload_index_arc.clone(),
                     quantized_vectors_arc,
                     Some(permit.clone()),
+                    gpu_device
+                        .as_ref()
+                        .map(|gpu_device| gpu_device.locked_device.clone()),
                     stopped,
                 )?;
             }

lib/segment/src/segment_constructor/segment_constructor_base.rs

@@ -364,6 +364,7 @@ pub(crate) fn create_vector_index(
     payload_index: Arc<AtomicRefCell<StructPayloadIndex>>,
     quantized_vectors: Arc<AtomicRefCell<Option<QuantizedVectors>>>,
     permit: Option<Arc<CpuPermit>>,
+    gpu_device: Option<Arc<gpu::Device>>,
     stopped: &AtomicBool,
 ) -> OperationResult<VectorIndexEnum> {
     let vector_index = match &vector_config.index {
@@ -379,6 +380,7 @@ pub(crate) fn create_vector_index(
                 payload_index,
                 hnsw_config: vector_hnsw_config.clone(),
                 permit,
+                gpu_device,
                 stopped,
             };
             if vector_hnsw_config.on_disk == Some(true) {
@@ -559,6 +561,7 @@ fn create_segment(
             payload_index.clone(),
             quantized_vectors.clone(),
             None,
+            None,
             stopped,
         )?);
 

lib/segment/src/vector_storage/quantized/quantized_vectors.rs

@@ -111,11 +111,11 @@ impl fmt::Debug for QuantizedVectorStorage {
 
 #[derive(Debug)]
 pub struct QuantizedVectors {
-    storage_impl: QuantizedVectorStorage,
-    config: QuantizedVectorsConfig,
-    path: PathBuf,
-    distance: Distance,
-    datatype: VectorStorageDatatype,
+    pub storage_impl: QuantizedVectorStorage,
+    pub config: QuantizedVectorsConfig,
+    pub path: PathBuf,
+    pub distance: Distance,
+    pub datatype: VectorStorageDatatype,
 }
 
 impl QuantizedVectors {

lib/segment/tests/integration/batch_search_test.rs

@@ -170,6 +170,7 @@ fn test_batch_and_single_request_equivalency() {
         payload_index: payload_index_ptr,
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/byte_storage_hnsw_test.rs

@@ -223,6 +223,7 @@ fn test_byte_storage_hnsw(
         payload_index: segment_byte.payload_index.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/byte_storage_quantization_test.rs

@@ -350,6 +350,7 @@ fn test_byte_storage_binary_quantization_hnsw(
         payload_index: segment_byte.payload_index.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/exact_search_test.rs

@@ -140,6 +140,7 @@ fn exact_search_test() {
         payload_index: payload_index_ptr.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/filtrable_hnsw_test.rs

@@ -200,6 +200,7 @@ fn _test_filterable_hnsw(
         payload_index: payload_index_ptr.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/hnsw_discover_test.rs

@@ -118,6 +118,7 @@ fn hnsw_discover_precision() {
         payload_index: payload_index_ptr,
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();
@@ -241,6 +242,7 @@ fn filtered_hnsw_discover_precision() {
         payload_index: payload_index_ptr,
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/hnsw_quantized_search_test.rs

@@ -141,6 +141,7 @@ fn hnsw_quantized_search_test(
         payload_index: segment.payload_index.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/multivector_filtrable_hnsw_test.rs

@@ -192,6 +192,7 @@ fn test_multi_filterable_hnsw(
         payload_index: payload_index_ptr,
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/multivector_hnsw_test.rs

@@ -150,6 +150,7 @@ fn test_single_multi_and_dense_hnsw_equivalency() {
         payload_index: segment.payload_index.clone(),
         hnsw_config: hnsw_config.clone(),
         permit: Some(permit.clone()),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();
@@ -164,6 +165,7 @@ fn test_single_multi_and_dense_hnsw_equivalency() {
         payload_index: segment.payload_index.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

lib/segment/tests/integration/multivector_quantization_test.rs

@@ -320,6 +320,7 @@ fn test_multivector_quantization_hnsw(
         payload_index: segment.payload_index.clone(),
         hnsw_config,
         permit: Some(permit),
+        gpu_device: None,
         stopped: &stopped,
     })
     .unwrap();

nvidia_icd.json

@@ -0,0 +1,7 @@
+{
+    "file_format_version" : "1.0.0",
+    "ICD": {
+        "library_path": "libGLX_nvidia.so.0",
+        "api_version" : "1.3"
+    }
+}

run.bat

@@ -0,0 +1,7 @@
+SET QDRANT__GPU__indexing=1
+SET QDRANT__GPU__max_groups=512
+SET QDRANT__GPU__force_half_precision=0
+
+SET QDRANT__storage__optimizers__max_optimization_threads=1
+
+cargo run --release

run.sh

@@ -0,0 +1,9 @@
+export QDRANT__GPU__indexing=1
+export QDRANT__GPU__max_groups=512
+export QDRANT__GPU__force_half_precision=0
+
+export QDRANT__GPU__device_filter=LLVM
+
+export QDRANT__storage__optimizers__max_optimization_threads=1
+
+cargo run --release

src/main.rs

@@ -168,6 +168,24 @@ fn main() -> anyhow::Result<()> {
             .async_scorer
             .unwrap_or_default(),
     );
+    if let Some(settings_gpu) = &settings.gpu {
+        segment::index::hnsw_index::gpu::set_gpu_indexing(settings_gpu.indexing);
+        segment::index::hnsw_index::gpu::set_gpu_force_half_precision(
+            settings_gpu.force_half_precision,
+        );
+        segment::index::hnsw_index::gpu::set_gpu_max_groups(settings_gpu.max_groups);
+        segment::index::hnsw_index::gpu::set_device_index(settings_gpu.device_index);
+        segment::index::hnsw_index::gpu::set_devices_count(settings_gpu.devices_count);
+        segment::index::hnsw_index::gpu::set_device_filter(&settings_gpu.device_filter);
+        segment::index::hnsw_index::gpu::set_wait_free(settings_gpu.wait_free.unwrap_or(true));
+        segment::index::hnsw_index::gpu::set_gpu_parallel_indexes(
+            settings_gpu.parallel_indexes.unwrap_or(1),
+        );
+        // initialize GPU devices
+        let _ = segment::index::hnsw_index::gpu::GPU_DEVICES_MANAGER
+            .as_ref()
+            .map(|devices_manager| devices_manager.lock_device());
+    }
 
     welcome(&settings);
 

src/settings.rs

@@ -141,6 +141,30 @@ pub struct TlsConfig {
     pub cert_ttl: Option<u64>,
 }
 
+#[derive(Clone, Debug, Deserialize, Validate)]
+#[allow(dead_code)]
+pub struct GpuConfig {
+    #[serde(default)]
+    pub indexing: bool,
+    #[serde(default)]
+    pub force_half_precision: bool,
+    #[serde(default)]
+    #[validate(range(min = 1))]
+    pub max_groups: Option<usize>,
+    #[serde(default)]
+    pub device_index: Option<usize>,
+    #[serde(default)]
+    pub devices_count: Option<usize>,
+    #[serde(default)]
+    pub device_filter: String,
+    #[serde(default)]
+    pub min_points: Option<usize>,
+    #[serde(default)]
+    pub wait_free: Option<bool>,
+    #[serde(default)]
+    pub parallel_indexes: Option<usize>,
+}
+
 #[derive(Debug, Deserialize, Clone, Validate)]
 #[allow(dead_code)] // necessary because some field are only used in main.rs
 pub struct Settings {
@@ -168,6 +192,9 @@ pub struct Settings {
     pub load_errors: Vec<LogMsg>,
     #[serde(default)]
     pub inference: Option<InferenceConfig>,
+    #[serde(default)]
+    #[validate(nested)]
+    pub gpu: Option<GpuConfig>,
 }
 
 impl Settings {