6月19日、Anubhab Banerjeeが「GPU-Resident Top-K for Agentic RAG: I Built a CUDA Kernel So My Retrieval Step Would Stop Bouncing Off the GPU」と題した記事を公開した。コーパスもクエリ埋め込みもすでにGPU上にあるのに、Agentic RAGの検索ステップはなぜかCPUで処理されている——この逆説に気づいたBanerjeeがCUDAカーネルを自作し、PCIeバス転送のオーバーヘッドを排除した結果、GTX 1080という7年前のGPUで最大8.57倍の高速化を達成した実装の記録だ。
問題の本質：GPUがあるのにCPUで検索している
Agentic RAGパイプラインでよく見られる実装では、エージェントがツールコールで知識ベースを検索するたびに、以下のことが起きている。
- クエリ埋め込みをGPUで生成する
- その埋め込みをホストRAMに転送する（PCIe経由）
- CPUでNコーパス行のドット積を計算し、Top-Kを選ぶ
- インデックスとスコアをGPUに転送して戻す
コーパスはすでにVRAMにある。クエリ埋め込みもGPU上で生成した。それにもかかわらず、検索のたびにデータをホストへ往復させている。1エージェントが1推論ステップで10回ツールコールするなら、このPCIe往復が10回発生する。コーパスサイズN=100万・次元数D=1024の規模では、スコアリングではなくこの往復コスト自体が検索レイテンシの支配的な要因になる。
解決策：コーパスをVRAMに常駐させる
アーキテクチャの骨格は単純だ。
agent.embed(query)
→ cudaMemcpy H→D (D floats)
→ row_dot_scores_kernel
→ partial_topk_block_kernel (P blocks)
→ merge_partial_topk_kernel
→ cudaMemcpy D→H (K indices + K scores)
コーパスは起動時に一度だけVRAMへアップロード（upload_corpus_rowmajor）し、以降はクエリごとにD次元の埋め込みを送り、2K個の結果（インデックスK個＋スコアK個）を受け取るだけだ。ホストとGPU間を行き来するデータ量がクエリあたり最小限に抑えられる。
CUDAカーネル実装の要点
Top-KをGPU上で実行する難しさ
ベクトルのスコアリング（行列積）はGPUが得意とする処理だ。問題は選択の部分にある。
全N行をソートしてTop-Kを得るO(N log N)のアプローチはコストが高すぎる。argpartitionに相当するツリー走査はGPUのメモリコアレッシング（連続メモリアクセスによる帯域効率化）を破壊する。thrustやcubなどのライブラリに頼ると、軽量なC++パイプラインに巨大なビルド依存が生じる。
本実装が採用したのは、意図的にシンプルな設計だ。各ブロックが担当する行範囲でシングルスレッドのO(K²)バブルソートを走らせ、最後にブロック間をシリアルにマージする。
__device__ void bubble_downward(float* const s, int* const ids, const int n) {
// K <= kMaxSupportedK の範囲でのみ動作する小規模ソート
for (int i = 0; i < n - 1; ++i) {
for (int j = 0; j < n - 1 - i; ++j) {
if (device_is_better(s[j + 1], ids[j + 1], s[j], ids[j])) {
const float ts = s[j];
s[j] = s[j + 1];
s[j + 1] = ts;
const int ti = ids[j];
ids[j] = ids[j + 1];
ids[j + 1] = ti;
}
}
}
}
V1カーネルの設計方針は「賢さより検証可能性」だ。343行のコードは短時間で通読でき、CPUオラクルと出力を1ビット単位で照合できる。
K≤32ではこのシングルレーンバブルソートで十分なパフォーマンスが出る。一方、K=100では性能が落ちる。これはバブルソートのO(K²)コストが原因だ。K=8なら各ブロックの選択ループは約28回の比較で済むが、K=100では約4,950回にまで膨れ上がる。GPU上の多数のスレッドが同じブロックの選択フェーズで直列的に待機することになり、GPU本来の並列性を活かせなくなる。記事中でこの限界は隠さず示されており、次バージョンではワープ特化のトーナメントセレクタへの置き換えが予告されている。
タイブレークの一致問題
スコアが浮動小数点精度で完全に一致する場合、CPUとGPUが異なる行を「勝者」と判定するとベンチマークが無意味になる。本実装ではホスト側とデバイス側のコンパレータを同一のロジックで実装している。
// ホスト側
bool is_better_score_pair(float32_t score_lhs, index_t idx_lhs,
float32_t score_rhs, index_t idx_rhs) {
if (score_lhs != score_rhs) return score_lhs > score_rhs;
return idx_lhs < idx_rhs; // スコア同一時は行インデックスが小さい方を優先
}
「スコアが高い方、同点なら行インデックスが小さい方」という一文で定義し、両側に同じ実装を置くことで、ベンチマークの不一致がロジックのバグによるものか確認できる体制を作っている。
ベンチマーク結果
GTX 1080（7年前のハードウェア）で、N∈{10k, 50k, 100k, 500k, 1M}・D∈{384, 768, 1024}の45構成を網羅したスイープを実施した。
| K値 | GPU勝利構成数 | 最大スピードアップ |
|---|---|---|
| K=8 | 15/15 | 8.57×（N=1M, D=1024） |
| K=32 | 13/15 | 7.76× |
| K=100 | 1/15 | CPUが14/15構成で勝利 |
K=8とK=32ではGPU常駐パスが圧倒的に有利だ。K=100でCPUが逆転するのは、前述のO(K²)バブルソートコストが臨界点を超えるためであり、記事中で正直に示されている。「魔法のカーネル」ではなく「コーパスをホストRAMに無駄に送るのをやめた」ことが勝因だとBanerjeeは強調する。
なお、K=8でGPUが負ける構成が存在しないのに対し、K=32ではN=10k・D=384とD=768の2構成でCPUが勝る。コーパスが小さい場合はPCIe往復コスト自体が安いため、GPUカーネルのオーバーヘッドが相対的に目立つ。
実装の公開
ソースコードはGitHubで公開されている（※編集部の考察：記事本文からリポジトリURLを特定できなかったため、直接リンクは省略する。元記事末尾にリンクが記載されている）。C++オーケストレータAPIはupload_corpus_rowmajor（起動時1回）とsearch_resident（クエリごと）の2関数で構成されており、既存パイプラインへの統合を想定した設計になっている。
Banerjeeは5G/6Gの基地局エンジニアリング出身で、基地局がビームコードブックの候補を受信電力でスコアリングして最良のK本を選ぶ処理がRAGのTop-Kと構造的に同一だという視点からこのカーネルを設計している。この記事は「Production-Grade Agentic Inference」シリーズの第3弾であり、第4弾ではエージェント間のハンドオフをまたいだ状態永続化（コールドスタート問題の解消）が扱われる予定だ。
関連技術として、FAISSやScaNNのようなGPU対応ベクトル検索ライブラリも存在するが、本実装はそれらへの依存を避け、Agentic RAGパイプラインへの最小統合コストを優先した点で位置づけが異なる（※編集部の考察）。
詳細はGPU-Resident Top-K for Agentic RAG: I Built a CUDA Kernel So My Retrieval Step Would Stop Bouncing Off the GPUを参照していただきたい。
