ベクトル検索を世界最速のものにするためにElasticsearch simdvecを構築した方法

Elasticsearchのすべてのベクトル検索クエリの基盤となる、手作業で調整されたSIMDカーネルライブラリElasticsearch simdvecの構築方法。

Elasticsearch simdvecは、Elasticsearch内のすべてのベクトル距離計算のエンジンです。Elasticsearchがサポートするすべてのベクトルタイプに手動で調整されたAVX-512とNEONカーネルを提供します。そのバルクスコアリングアーキテクチャは、x86では明示的なプリフェッチ、ARMではインターリーブロードによってメモリレイテンシを隠蔽し、データがCPUキャッシュを超える場合、FAISSやjvectorなどのライブラリを最大4倍上回るパフォーマンスを発揮します。この記事では、これを構築した理由、その内部構造、そしてこれがElasticsearchのベクトル検索を世界最速の1つにしている理由について説明します。

Elasticsearch simdvecの構築方法

Elasticsearchにおけるすべてのベクトル検索クエリは、Hierarchical Navigable Small World(HNSW)トラバーサル、転置ファイル(IVF)スキャン、再ランキングパスのいずれであっても、結局は同じ問題に帰着します。すなわち、クエリごとにベクトル間の距離を何百万回も計算しなければならないということです。Elasticsearchは、float32からint8、bfloat16、バイナリ、Better Binary Quantization(BBQ)まで、幅広いデータ型と量子化戦略をサポートしています。それぞれにメモリ、スループット、リコールのトレードオフが異なります。そのすべての背後にあるのは、simdvecという単一のエンジンです。

私たちは、ハードウェアが許す限りすべての距離計算を高速に行うためにsimdvecを構築しました。この記事では、simdvecを構築した理由、その内部構造、そしてどのような場面で最も効果を発揮するのかを説明します。

レースカーのような構造

F1愛好家として(当社チームには以前フェラーリのF1チームで勤務経験があるメンバーもいます)明確な類似点を見いだせます。F1カーは、最高のラップタイムを達成するという唯一の目的のために設計されています。エンジンのパワー、空力性能、シャーシ設計は、その結果に貢献する限りにおいてのみ重要です。ベクトルデータベースについても同様で、インデキシングのスループット、クエリのレイテンシ、リコールが成功を定義します。

最終結果は重要ですが、最高レベルのパフォーマンスを達成するには、各コンポーネントが最高の状態である必要があります。必要十分ではなく、そのカテゴリーで最高でなければなりません。simdvecはその考え方で構築されており、システムの重要な部分であるエンジンに焦点を当てています。これは、専用に構築された、単一命令多重データ(SIMD)最適化カーネルライブラリで、JavaからPanama外部関数インターフェース(FFI)を介して呼び出される、手動で調整されたネイティブC++距離関数を提供します。一括スコアリング、キャッシュラインのプリフェッチ、Elasticsearchで使用されるすべてのベクトルの種類とレイアウトをサポートしています。

それがすべてのクエリの背後にあるエンジンです。

自社開発した理由

当社は2023年にApache LuceneのPanama Vector APIを使用してスタートしました。float32の内積計算にはうまく機能しましたが、すぐにその機能ではElasticsearchのニーズには対応できなくなりました。Elasticsearchは、int8、int4、bfloat16、シングルビット、非対称BBQなど、幅広い量子化されたベクトルタイプをサポートしています。それぞれにSIMD戦略、パッキングレイアウト、アキュムレータの要件が異なります。型カバレッジに加えて、Elasticsearchのスコアリングパスは単一ペアのスループット以上のものを要求します。HNSWは1回のパスで複数のグラフ隣接ノードをスコアリングする必要があり、IVFはプリフェッチを使用して数千の候補を一括スコアリングする必要があり、ディスクベースのスコアリングはコピーせずにmmapされたメモリ上で直接動作する必要があります。入手可能なものを調べてみましたが、すべての条件を満たすものは見つかりませんでした。

そこで、simdvecを構築しました。これは、JavaからFFI経由で呼び出される手作業で調整されたネイティブC++カーネルで、一括スコアリング、プリフェッチ、Elasticsearchが使用するすべてのベクトルタイプをサポートしています。ライブラリを所有することで、私たちはフルスタックを制御できます。BBQのような新しい量子化タイプを追加すると、システム全体にわたって調整されたSIMDカーネルが組み込まれます。上流のライブラリがそれをサポートするのを待つ必要はなく、いかなる型においてもパフォーマンスに妥協することはありません。Elasticsearchにおけるすべてのベクトルクエリ(HNSW、IVF、リランキング、ハイブリッドなど)は、実際に使用する操作とタイプに基づいて構築されたこのエンジン上で実行されます。

simdvecには、x86とARMそれぞれに対応したネイティブライブラリが用意されており、起動時に複数の命令セットアーキテクチャ(ISA)階層を選択できます。FFIを介したJavaからの呼び出しオーバーヘッドは非常に低く、1桁ナノ秒です。

ランドスケープ

SIMD最適化されたベクトル距離カーネルを構築しているのは私たちだけではありません。このエコシステムは豊かで、私たちはsimdvecがどのように動作するのかを理解したいと考えました。プロジェクトの優劣を決めるのではなく、背景情報を提供し、Elasticsearchエンジンの位置づけを説明することが目的です。私たちは異なるアプローチを表す3つのプロジェクトを基準点として選びました。

  • jvector:ベクトル化された距離計算にPanama Vector APIを使用するJava近似最近傍探索(ANN)ライブラリ。x86上ではオプションでネイティブCアクセラレーションが可能。
  • FAISS:手作業で調整されたAVX2/AVX-512カーネルを備えた、広く展開されているオープンソースのベクトル検索フレームワーク。
  • NumKong (旧称 SimSIMD):距離関数、行列演算、地理空間計算など、2,000種類以上の手作業で調整されたSIMDカーネルを網羅した包括的なスイート。

各プロジェクトは異なる目的を持ち、異なるトレードオフをもたらします。Elasticsearchが必要とする特定の操作におけるsimdvecのパフォーマンスのコンテキストを提供するために、それらからの参照番号を含めています。

測定方法

simdvecとjvectorのベンチマークは、FFIオーバーヘッドを含めた標準JVMマイクロベンチマークハーネスであるJMHを使用してJavaで記述されています。NumKongベンチマークFAISSベンチマークについては、標準的なC++マイクロベンチマークフレームワークであるGoogle Benchmarkを使って小さなC/C++ハーネスを作成しました。どちらのフレームワークも、ウォームアップと反復キャリブレーションを含めた1操作あたりのナノ秒数の精度を報告しています。ハードウェアパフォーマンスカウンターを介して、すべてのライブラリが両方のプラットフォームでSIMDを使用していることを確認しました。ベンチマークコードはすべて、リンク先のGitHubリポジトリ(およびsimdvecの場合はelasticsearchリポジトリ)で公開されています。

ソフトウェア: JDK 25.0.2、JMH 1.37、GCC 14、Google Benchmark(最新版)。

一度に1つのベクトル

ベクトル検索における最も基本的な操作は、2つのベクトル間の距離を計算することです。すべてのHNSW近隣評価、すべてのIVF候補スコア、すべての再ランク比較は、この内側のループに還元されます。

両プラットフォームで1024次元における単一ペアのスループットを測定しました。まず、ベースラインとなる型であり、エコシステムの競争が最も激しいfloat32から評価を開始しました。simdvecをFAISSおよびjvectorと比較しました。NumKongはfloat32にfloat64アキュムレータを使用するため、スループットよりも数値精度を優先し、処理速度が3.2倍から5.3倍遅くなる(プラットフォームによって異なる)ため、比較対象から除外しました。比較を同じように保つために、代わりにint8でNumKongのベンチマークを行います。ここでは、simdvecと同じアキュムレータ戦略を使用しています。

x86アーキテクチャーでは、FAISS AVX-512が23ナノ秒と最速のシングルペアカーネルです。simdvec AVX-512は28ナノ秒で続きますが、この差はFFI呼び出しのオーバーヘッドを反映したものです。どちらもマルチアキュムレータアンローリングを備えた512ビットFMAを使用しています。AVX2レベルでは、両者の性能差ははるかに小さく、それぞれ36ナノ秒と39ナノ秒で、どちらも256ビットのレジスタとメモリのロード幅によって制約されています。jvectorはJava Panama Vector APIを使用して44ナノ秒で到達します。Panamaは優れたSIMDコードを生成しますが、手作業で調整されたC++ intrinsicsが優位を保っています。

ARMでは、simdvecは70ナノ秒でリードしており、110ナノ秒のjvectorと156ナノ秒のFAISSをはるかに上回っています。simdvecはaarch64向けにNEONカーネルを独自にチューニングしました。JvectorにはネイティブのARMコードがなく、Panamaに依存しています。FAISSは明示的なNEON組み込み関数ではなく、コンパイラの自動ベクトル化に依存しているため、差が大きくなっています。これは、カーネルライブラリを所有することによる実用的な利点を反映しています。ElasticsearchがGravitonに拡張された際、専用に構築されたNEONカーネルを追加しました。jvectorもFAISSも、ARMネイティブコードを同程度に優先しているわけではありません。

しかし、Elasticsearchはfloat32だけを評価するわけではありません。Int8の量子化はメモリ使用量を4分の1に削減し、bfloat16は2分の1に、BBQは32分の1に削減します。それぞれのタイプには独自のSIMD戦略が必要であり、simdvecはすべてのタイプ向けに手動チューニングされたネイティブカーネルを提供しています。

比較したライブラリの中で、int8用の同等のカーネルを備えているものはNumKongのみでした。int8ドット積、二乗ユークリッド、余弦を1024次元で測定しました。

Int8シングルペアスコアリング(1024次元、ns/vec 演算 - 数値が小さいほど良い)

どちらのアーキテクチャでも、NumKongは小~中次元では同等かより高速です。その差は主に呼び出しのオーバーヘッドが少ないことに起因します(Cの直接呼び出しとJavaのFFI)。より大きな次元では、simdvecが追いつき、より効率的なカーネル実装(カスケードアンローリングを使用する)が呼び出しコストを償却します。次元が増えるにつれて、このギャップは閉じ、最終的には逆転します。クロスオーバーのサイズは、機能や構造によって768から1536の間となります。

Java FFIのオーバーヘッドが若干高いにもかかわらず、simdvecは高度に最適化されたC/C++ライブラリと同等の性能を発揮します。float32int8の両方に対応した最適化されたカーネルを備えた唯一のライブラリであるだけでなく、ARMアーキテクチャではトップクラスの性能を誇り、x86アーキテクチャではFAISSにわずかに劣る程度(float32の場合)、そして両アーキテクチャにおいてNumKongに非常に近い性能(int8の場合)を実現しています。また、bfloat16、int4、binary、BBQについては、代替手段は存在するものの、simdvecは各型のデータレイアウトに合わせて手作業で調整されたSIMDによって差別化を図っています。

しかし、実際の検索エンジンは一度に1つのベクトルを評価するのではなく、クエリごとに数千のベクトルを評価します。次の質問は、そのスケールで何が起こるかということです。

一度に数千件

シングルペアのパフォーマンスは全体の一部に過ぎません。実際には、システムが負荷時にどのように動作するかが重要です。単一のHNSWクエリは数百のグラフ近傍をスコアリングすることがあります。IVFスキャンは、数千件の投稿リストエントリーをスコアリングすることがあります。リランクパスは数万の候補をスコアリングすることがあります。シングルペアのスループットは重要ですが、より重要なのは、多くのベクトルをどれだけ速くスコアリングできるか、そして作業セットがCPUキャッシュから溢れ出すにつれてパフォーマンスがどれだけスムーズに低下するかです。

simdvecは、あらゆるデータタイプに対して一括スコアリング機能を提供します。これらは単なる単一ペアカーネル上のループではなく、クエリベクトルを1次元ストライドごとに1回ロードし、複数のドキュメントベクトル間で共有するマルチアキュムレータ内部ループを使用します。次のバッチに対しては、明示的なキャッシュラインプリフェッチが行われます。jvectorもFAISSも同等の機能を提供していません(執筆時点では)。Jvectorには一括処理APIがないため、呼び出し元はループ内で一度に1組ずつスコアを計算します。FAISSはfvec_inner_products_nyを公開していますが、執筆時点では、クエリ償却やプリフェッチを行わずに、シングルペア距離関数のループとして実装されています。

Float32。カーネルレベルでのインパクトを測定するために、HNSWのような散在グラフの近傍検索をシミュレートするランダムアクセスパターンを用いて、1024次元float32ドキュメントベクトルの数を増やしながら単一のクエリのスコアを計算しました。3つのデータセットサイズ(32、625、32,500ベクトル)は、それぞれL1、L2、L3キャッシュを超えるように選択されています。

データがキャッシュに収まる場合、simdvecはどちらのプラットフォームでも最速ですが、カーネル演算が支配的であるため、マージンは控えめです。実際の分離は、ワーキングセットがL3を超えるにつれて現れます。x86では、simdvecのスコアは1ベクトルあたり95ナノ秒ですが、FAISSは165ナノ秒、jvectorは412ナノ秒です。ARM上でも同様の傾向が見られ、simdvecは162ナノ秒で安定しているのに対し、FAISSは347ナノ秒、jvectorは476ナノ秒に上昇します。simdvecのプリフェッチとクエリの償却により、シングルペアカーネル上の単純なループでは対応できない方法でメモリの待ち時間が隠され、メインメモリの奥深くで実際の検索ワークロードが動作する場所でその利点が広がります。

Int8。同じパターンは量子化された型にも当てはまります。int8ドット積の一括スコアリングを1024次元で測定し、データセットのサイズが同じL1、L2、L3キャッシュ境界を超えるように選択して、simdvecのバルクスコアリングをループ内のNumKongシングルペアスコアリングと比較しました。

x86では、simdvecは1.2倍~1.9倍高速で、これは明示的なプリフェッチとバッチ処理の組み合わせによってもたらされます。ARM環境では、すべてのデータセットサイズにおいて、simdvecが再び優位に立ちました(1.7倍から1.9倍高速)。その利点は、4つのベクトルを一度にバッチ処理することで、インターリーブアクセスパターンを介してメモリレベルの並列処理を実現することにあります。いずれの場合も、最も顕著な結果は最大のデータセットサイズで何が起こるかであり、それが最も重要な場所です。

結果は、二乗距離とコサインで同様のパターンを示し、ARMでは1.4倍から1.8倍、x86では1.3倍から3.0倍の速度向上が見られました(詳細はこちら)。

メモリが重要となる場合

本番環境のベクトルインデックスは通常、CPUキャッシュに収まりません。1024次元の10Mベクトルint8インデックスは10GBです。候補のスコアリングとは、DRAMからデータをストリーミングすることを意味し、そこでバルクスコアリングアーキテクチャが大きな違いを生むのです。

バルクスコアリング中のCPU内部で何が起こるかを測定するためにハードウェアパフォーマンスカウンターを使用し、メモリーレイテンシを隠すためには、アーキテクチャごとに根本的に異なる2つの戦略が必要であることを発見しました。

x86アーキテクチャでは、明示的なプリフェッチによってキャッシュミスが解消されます。バルクカーネルはベクトルを逐次的に処理し、次のベクトルを処理する前に1つのベクトルを完全に計算しながら、次のバッチのためのプリフェッチ命令を発行します。将来のデータはCPUが必要とする前にL1に取り込まれます。

ARMアーキテクチャでは、プリフェッチを使用しても、同じ逐次的なアプローチではパフォーマンスが低下しました。代わりに、バルクカーネルが4つのベクターからすべてのストライド位置でロードをインターリーブし、アウトオブオーダーエンジンに4つの独立したメモリストリームを提供します。CPUのデータの取得速度が速くなったわけではなく、メモリ要求の処理中に常に別の計算処理を行うことで、待ち時間を短縮しているのです。詳細な分析についてはこちらのGitHubイシューをご覧ください。

数字は2つの異なる物語を語っています。

  1. x86アーキテクチャでは、プリフェッチによって139,000回のキャッシュミスが19,000回に減り、1サイクルあたりの命令実行数(IPC)が2倍以上になります。データセットのサイズが大きくなるにつれてプリフェッチによってコストのかかるDRAM往復処理が徐々に隠蔽されるため、データ量の増加によるメリットは大きくなり、L2レベルでは1.2倍、L3レベルを超えると2.8倍になります。
  2. ARMではキャッシュミスはほとんど変わりません。変化するのは利用率です。インターリーブアクセスパターンによってパイプラインへの供給が維持されるため、バックエンドの停止時間が40%減少します。この利点は、データセットのサイズに関係なく一貫して1.8倍です。これは、データがキャッシュまたはDRAMから来ているかに関係なく、メモリレベルの並列処理が適用されるためです。

2つのアーキテクチャ、2つの戦略、結果は1つ:本番環境規模では、simdvecはベクトルがメインメモリ全体に分散している場合でも、CPUパイプラインを常にフル稼働させます。

Elasticsearchユーザーへの影響

これらのカーネルレベルの能力は相乗効果を発揮します。単一のベクトル検索クエリでは、HNSWグラフの走査、候補のスコアリング、再ランキングなど、数百万もの距離演算が計算される場合があります。数千回の同時クエリにおいて、1回の操作でナノ秒単位がクエリ遅延やクラスタスループットに直接変換されます。float32、int8、bfloat16、BBQのいずれを使用する場合でも、インデックスがメモリ上にあるかディスク上にあるかに関わらず、simdvecが基盤となるエンジンであり、これらのすべての操作は同じエンジンを通して実行され、ナノ秒単位まで最適化されています。

重要なポイントは、本番規模では、ベクトル検索のパフォーマンスは主に生のSIMDスループットによって決まるわけではないということです。システムのパフォーマンスは、数百万もの小さな演算処理において計算能力を維持しながら、メモリの遅延をいかに効率的に隠蔽できるかに大きく左右されます。

simdvecカーネルは、ほぼすべてのElasticsearchリリースで改善されています。新しい量子化タイプやハードウェアプラットフォームが登場すると、初日からチューニング済みのカーネルが提供されます。また、既に出荷されている実装を改良していくにつれて、既存の型もますます高速化していきます。

このコンテンツはどれほど役に立ちましたか?

役に立たない

やや役に立つ

非常に役に立つ

関連記事

最先端の検索体験を構築する準備はできましたか?

十分に高度な検索は 1 人の努力だけでは実現できません。Elasticsearch は、データ サイエンティスト、ML オペレーター、エンジニアなど、あなたと同じように検索に情熱を傾ける多くの人々によって支えられています。ぜひつながり、協力して、希望する結果が得られる魔法の検索エクスペリエンスを構築しましょう。

はじめましょう