WebGPUでベクトル検索を
実装する

RAG（Retrieval-Augmented Generation）は通常、サーバー上のベクトルDBとLLMで構成します。 しかしすべてブラウザ内で完結させることは可能でしょうか。 この記事では、RAGの心臓部であるベクトル検索だけを取り出し、 WebGPU で実装できるかを検証します。

ベクトル検索の核心は、「意味が近いテキストは、ベクトル空間でも近くに配置される」という性質です。 埋め込みモデルは、テキストを384次元の数値ベクトルに変換します。

// ESモジュールとしてCDNから直接import
import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.0.0';

// 埋め込みパイプラインを作成（WebGPU優先、非対応ならWASM）
const extractor = await pipeline(
  'feature-extraction',
  'Xenova/multilingual-e5-small',   // 日本語対応・384次元
  { device: 'webgpu' }                    // ← これだけでGPU実行
);

// テキストを384次元ベクトルに変換（L2正規化つき）
const output = await extractor('WGSLのバリア同期について', {
  pooling: 'mean',      // トークン平均を取る
  normalize: true,      // L2正規化 → 内積=コサイン類似度に
});
const vec = output.data;  // Float32Array(384)

device: 'webgpu' を指定するだけで、ONNX Runtime Web が 埋め込みモデルの推論をGPUで実行します。非対応環境では自動的にWASM（CPU）にフォールバックするため、 エラーハンドリングは不要です。

// 1スレッド = 1文書の類似度を担当
@group(0) @binding(0) var<storage, read>       query: array<f32>;   // 384
@group(0) @binding(1) var<storage, read>       docs:  array<f32>;   // N×384
@group(0) @binding(2) var<storage, read_write> scores: array<f32>;  // N
@group(0) @binding(3) var<uniform> p: Params;  // numDocs, dim

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3u) {
  let docId = gid.x;
  if (docId >= p.numDocs) { return; }

  // この文書のベクトルとクエリの内積を計算
  var dot = 0.0;
  let offset = docId * p.dim;   // 文書ベクトルの開始位置
  for (var i = 0u; i < p.dim; i = i + 1u) {
    dot = dot + query[i] * docs[offset + i];
  }
  // 正規化済みなので、内積 = コサイン類似度
  scores[docId] = dot;
}

ハイライト部分が検索の核心です。各スレッドが1文書を担当し、384次元の内積を計算します。 1万文書なら1万スレッドが同時に走ります。行列積（記事⑪）の特殊ケース—— クエリが1行だけの「行列 × ベクトル」演算と見ることもできます。

// GPUから類似度スコアを取得
const scores = new Float32Array(readBuf.getMappedRange());

// スコアとインデックスをペアにしてソート
const ranked = Array.from(scores)
  .map((score, idx) => ({ score, idx }))
  .sort((a, b) => b.score - a.score)   // 降順
  .slice(0, topK);                      // 上位K件

// 文書本体と紐付けて返す
return ranked.map(r => ({
  ...knowledgeBase[r.idx],
  score: r.score,
}));

ソート自体はCPUで行います。文書数が数万程度なら一瞬です。 百万件を超えるとGPU上でのソート（bitonic sort等）が必要になりますが、 ブラウザで扱う規模では過剰最適化です。

ベクトル検索は本当にGPUで速くなるのでしょうか。 合成ベクトルで文書数を増やしながら、CPU実装と比較します。

Android（Pixel 10）で実行して、設計レベルの重大なバグを2つ発見しました。 どちらも実機で動かさなければ気づけないものでした。

// ❌ 悪い実装：検索のたびに全文書を転送（73MB × 毎回）
async function searchGPU(queryVec, docVecs, numDocs) {
  const dBuf = device.createBuffer({ size: numDocs*DIM*4, ... });
  device.queue.writeBuffer(dBuf, 0, docVecs);  // ← 毎回73MB！
  ...
}

// ✅ 正しい実装：文書は一度だけ転送して常駐
function createSearchContext(docVecs, numDocs) {
  const dBuf = device.createBuffer({ size: numDocs*DIM*4, ... });
  device.queue.writeBuffer(dBuf, 0, docVecs);  // ← 起動時に1回だけ
  return { qBuf, dBuf, sBuf, uBuf, rBuf, bind, numDocs };
}

async function searchWithCtx(ctx, queryVec) {
  device.queue.writeBuffer(ctx.qBuf, 0, queryVec);  // ← クエリ1.5KBのみ
  // ディスパッチしてスコアを読むだけ
}

GPUは「隣り合うスレッドが、隣り合うメモリを読む」ときに最高効率を発揮します。 これをコアレッシング（合体アクセス）と呼びます。 複数のメモリ要求を1本のトランザクションにまとめられるためです。

// ❌ ナイーブ：文書ごとに連続（row-major）
for (var i = 0u; i < p.dim; i = i + 1u) {
  dot = dot + query[i] * docs[docId * p.dim + i];
  //                          ↑ スレッドごとに384離れる
}

// ✅ コアレス：次元ごとに連続（転置 / dimension-major）
for (var i = 0u; i < p.dim; i = i + 1u) {
  dot = dot + query[i] * docsT[i * p.numDocs + docId];
  //                            ↑ 隣接スレッドが隣接アドレス
}

// 転置はアップロード時に一度だけ行えばよい（CPU側で実行）
function transposeDocs(docVecs, numDocs) {
  const out = new Float32Array(numDocs * DIM);
  for (let d = 0; d < numDocs; d++)
    for (let i = 0; i < DIM; i++)
      out[i * numDocs + d] = docVecs[d * DIM + i];
  return out;
}

さらに一歩：vec4f で1スレッドが4文書を処理

コアレス化してもまだ帯域を使い切れていませんでした（実効7 GB/s、モバイルGPUは25〜30 GB/s級）。 原因はスカラー f32 を1個ずつ読んでいたこと。GPUは16バイト単位で読む方が効率的です。

// docsT4[i * numDocs4 + d4] は vec4f
// = 次元 i における「4文書分」の値がまとまったもの
@group(0) @binding(1) var<storage, read> docsT4: array<vec4<f32>>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3u) {
  let d4 = gid.x;  // 4文書ぶんのインデックス
  if (d4 >= p.numDocs4) { return; }

  var acc = vec4<f32>(0.0, 0.0, 0.0, 0.0);
  for (var i = 0u; i < p.dim; i = i + 1u) {
    // スカラーを vec4 にブロードキャスト乗算 → 4文書を一度に
    acc = acc + query[i] * docsT4[i * p.numDocs4 + d4];
  }
  scores4[d4] = acc;  // 4件のスコアを一度に書き出す
}

メモリ命令が1/4になり、1命令で16バイトを読むためコアレスも完璧になります。 文書数は4の倍数にパディングし、余りは0埋め（スコア0になるので無害）します。

Pixel 10 と iPad mini で同じベンチマークを実行したところ、 最速のカーネルが逆転しました。