はじめに

ローカル LLM 基盤を OpenAI 互換の入口でまとめる方針はすでに見えていたものの、埋め込みと rerank をどこまで分離して設計するかはまだ曖昧でした。今回のメモでは、その中でも検索基盤の再利用性に直結する embed と rerank の仕様だけを切り出して、Rust + Axum + ort + tokenizers で実装する前提を先に固定しました。

やりたかったのは、単に embedding API を生やすことではありません。文書側の計算を事前に済ませ、LAN 内で key -> bin キャッシュを共有しながら、検索段階と再ランキング段階を無駄なくつなぐ構成を先に決めておくことでした。関連メモでまとめていた LiteLLM Proxy や CPU/GPU 分業の発想とも整合するように、今回はあえて ONNX/CPU に寄せて境界をはっきりさせています。

背景・動機

既存のローカル LLM パイプライン案では、CPU 側が前処理、embedding、タスク整形を担当し、GPU 側が最終生成を引き受ける構成を考えていました。その流れで見ると、embedding と rerank は「生成の前段」にある独立した共通部品です。ここを曖昧なまま実装すると、後から RAG、候補検索、再順位付け、キャッシュ、gRPC 境界の設計が全部ぶれます。

特に rerank は、単一ベクトルの検索 API と違って token-level のベクトルを扱うので、I/O の粒度を最初から決めておかないと破綻しやすいです。文書 token vectors を毎回送るのは重いし、検索サービスと rerank サービスの間で大きな配列を投げ始めると、LAN 内でも帯域と待ち時間が効いてきます。だから今回は、「transport は key と最小メタだけ」「本文書の token vectors はキャッシュから引く」という方針を明示しました。

設計方針

まず方針として、仕様を embed(256d) と rerank(ColBERT 64d token, MaxSim) の二系統に明確に分けました。両方とも ONNX/CPU 前提に寄せて、実装言語も Rust に統一しています。関連メモでは LiteLLM や Task Router を含む全体パイプラインを考えていましたが、このノートではその内部にある「検索と再ランキングの基礎部品」を先に固める位置づけです。

ここで効いた判断は、doc 側のベクトルは事前計算、query 側はリクエスト時に計算するという役割分担でした。これで検索リクエストのレイテンシを抑えつつ、文書更新時だけ再計算すればよい形にできます。さらに key ベースのキャッシュを中心にすることで、同じ文書を複数ワーカーや複数サービスで使い回しやすくなります。

仕様詳細

設計目標

この設計での目標は明快です。embed 側は近傍検索用の単一ベクトルを返すこと、rerank 側は ColBERT の late interaction を使って候補の順序だけを精密化することです。ここで両者を混ぜないことで、検索の高速化と rerank の高精度化をそれぞれ独立に詰めやすくなります。

Embed API

POST /embed は texts: string[] を受け取り、256 次元のベクトルを返す API としました。処理は tokenizers で tokenize し、ort で ONNX 推論し、モデルの作法に合わせて pooling を行い、最後に 256 次元へ truncate して L2 normalize します。

ここで truncate と normalize を明示したのは、検索側の計算を安定させたいからです。元モデルの出力が 1024 次元級でも、そのまま返すより 256 次元で揃えた方が索引構造を設計しやすいし、cosine/inner product も扱いやすくなります。ただし、この部分はモデルカードの推奨とズレると品質が落ちるので、あとで Python 参照実装と数値一致を取る前提にしています。

Rerank API

POST /rerank は query: string と candidates: Candidate[] を受け取ります。Candidate は doc_key だけでもよいし、doc_id + seq + chunk_id を材料にしてその場で key を解決してもよい形です。返すのは scores: float32[] と order: int[] です。

rerank 側の肝は、query token vectors (Tq x 64) はその場で生成し、document token vectors (Td x 64) はキャッシュから引く点です。各 query token に対して document token 全体の最大類似度を取り、その合計でスコア化する MaxSim を採用しました。ColBERT の構造を保ちつつ、通信量を増やさない実装に寄せたかったので、この API では「候補全文を毎回送る」発想を避けています。

モデルとファイル構成

embed モデルは lightonai/modernbert-embed-large/onnx/* を前提にしました。最終的に INT8 (onnx/model_int8.onnx, 約 17MB) を採用しています。エッジ向けの粗さがむしろ検索精度に寄与し、レイテンシも約 10ms に収まりました。並列実行耐性を限界まで高める設計にしており、軽量な INT8 モデルはその方針にも合致しています。

tokenizer は同じリポジトリの tokenizer.json を使います。このとき special token や normalization 挙動を参照実装と揃えることが重要で、ここがズレると埋め込み品質の差がそのまま検索結果に出ます。

rerank モデルは mixedbread-ai/mxbai-edge-colbert-v0-32m を前提にしましたが、ここには注意点があります。もしリポジトリに ONNX が含まれていなければ、Rust から ort で直接回せません。そのため、実運用では lightonai/mxbai-edge-colbert-v0-32m-onnx を追加で取得する案を推奨にしています。

キャッシュ設計（LAN 内、key/bin 方式）

キャッシュは doc_id + content_seq を起点にした不可逆キーを使う方針にしました。材料は次の通りです。

  doc_id|content_seq|chunk_id|model_id|tokenizer_hash|max_len|chunk_ver|dtype|layout

これを blake3 でまとめて doc: プレフィックス付きのキーにします。重要なのは content_seq と chunk_ver を入れることでした。文書更新時だけ版を進める前提にしておけば、内容が変わらない限りキャッシュはそのまま再利用できますし、chunking 戦略を変えたときも chunk_ver で強制的に切り替えられます。

値のフォーマットは、まず fp32（per-row float32 vectors）で実装し、その後に bitpack（sign encoding + u64 array）を考える二段階にしました。最初から bitpack をやりたくなるのですが、実装と検証のしやすさを考えると、最初は fp32 で形を作った方が前に進みやすいです。

ヘッダには dtype, dim (=64), n_tokens, scales_present, layout (token-major), version を固定で入れます。token-major を明示したのは、MaxSim 実装とキャッシュデコードの両方でレイアウトの解釈を固定したかったからです。

gRPC 設計（バイナリ vec 直接送信）

当初は doc_keys のみを運搬し、vec はキャッシュから引く設計を考えていました。しかし実装段階で DB 依存を排除する方針に転換し、embed 済みの 256d vec を直接送信するアンチパターンを採用しました。

実装では gRPC の MaxSimSearch RPC を定義し、map<string, VecF32> で key と 256d ベクトルのペアを候補として受け取ります。256 floats × 4 bytes = 1024 bytes/vec で、1 候補あたり約 1KB です。加えて threshold、top_k、top_p によるフィルタリングをサーバ側で行い、結果は keys + scores の並列配列で返します。

  message MaxSimRequest {
  map<string, VecF32> candidates = 1;  // key -> 256d vec (100-200 typical)
  VecF32 query_vec = 2;               // query vector (256d)
  float threshold = 3;                // min score to include
  int32 top_k = 4;                    // max results
  float top_p = 5;                    // cumulative-score cutoff
}

一般的にはベクトルを RPC に乗せるのは避けるべきですが、以下の条件が揃っていたため合理的と判断しています。

  • ローカル LAN 内の通信に限定される
  • 1 候補あたり約 1KB、100-200 候補でも 100-200KB 程度に収まる
  • DB 依存がなくなることで、embed/rerank サービスが完全にステートレスになる
  • この精度で検索品質は十分という検証結果がある

結果として、キャッシュ層と PostgreSQL の content_seq 管理が不要になり、サービスの並列実行耐性を限界まで高める設計が可能になりました。

HTTP 側の /v1/rerank は別途テキストベースのインターフェースを持ち、query テキスト + candidate keys を受け取って内部で ONNX 推論を行います。gRPC はすでに embed 済みのベクトルを持つ上流サービス（agent-gateway）が MaxSim スコアリングだけを依頼するための高速パスです。

Rust 実装要件

Rust 実装では Axum を外向けの REST とし、必要なら内部だけ gRPC を使う前提にしました。tokenizers では tokenizer.json のロードと special token の一致、ort では intra_threads = num_cpus と optimization_level = Level3、そして最初は CPU execution provider 固定です。

MaxSim の類似度は内積を基本にしつつ、あとで bitpack + Hamming approximation に差し替えられるよう trait で切れるようにしておく想定です。ここを最初から抽象化しておくと、高速化の検証を本体ロジックから分離しやすくなります。

実装順序

実装順もノートの中で先に決めました。先に embed を完成させ、pool -> truncate256 -> normalize を固定する。その後に rerank の ONNX 入手可否を確認し、必要なら ONNX 配布版を追加で取得して /rerank を作る。さらに PostgreSQL の content_seq 更新ルールを入れて、doc 事前計算から cache set までのパイプラインを組み、最後に gRPC を最小構成で足します。

この順番にしたのは、先に外部境界を固定してから、後で transport を足す方がリスクが小さいからです。最初に gRPC から入ると、API、キャッシュ、配列形式の全部が同時に揺れます。

注意事項・未知数

一番大きい未知数は rerank モデルの ONNX 配布有無です。ここが欠けていると Rust 側の ort 実行は止まります。次に、embed の pooling と normalize をどの作法で固定するかも品質に直結します。さらに、chunking の max_len や stride を変えると全キャッシュが無効になるので、ここは最初から chunk_ver をキーに含めて事故を防ぐ必要があります。

結果

今回のメモで良かったのは、embedding と rerank の責務がかなり整理できたことでした。OpenAI 互換 Proxy 全体の中に埋め込むとしても、この 2 本は単独のサービスとして切り出せるし、逆に同一プロセスにまとめる判断もあとからできます。

当初は doc_key だけを運ぶ設計で、キャッシュと PostgreSQL の版管理を一本の線でつなぐ構想でした。しかし実装段階で DB 依存を排除する方針に転換し、embed 済みの 256d vec を gRPC で直接送信する形に落ち着きました。1 候補あたり約 1KB、100-200 候補でも 100-200KB 程度に収まるため、ローカル LAN 内であれば合理的な判断でした。結果として、キャッシュ層と PostgreSQL の content_seq 管理が不要になり、サービスの並列実行耐性を限界まで高める設計が可能になりました。

今後の作業

embed 側は INT8 モデル（約 17MB）で約 10ms のレイテンシに収まり、並列実行耐性も十分な水準に達しています。rerank 側も lightonai/mxbai-edge-colbert-v0-32m-onnx の ONNX 配布版を採用することで、Rust から ort で直接駆動できる状態になりました。

残っている設計ポイントは、/embed と /rerank を外部公開 API にするのか、OpenAI 互換 Proxy の内部 API に閉じるのかという境界の決定です。認証、監視、ワークフローの扱いがこの判断で変わります。

ただ、今回の時点で「何を RPC で運び、何を Rust 側で責務に持つか」は十分固定できました。DB 依存を排除してステートレスに振り切った判断は、実装を進める中で正解だったと感じています。