# TeleRAG: Efficient Retrieval-Augmented Generation Inference with Lookahead Retrieval > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 2 (May 19 / Tue)、15:15 - 15:30 PDT > - **登壇者:** Chien-Yu Lin ら / University of Washington, Harvard University, Shanghai Jiao Tong University > - **URL:** https://mlsys.org/virtual/2026/oral/3796 > - **OpenReview:** https://openreview.net/forum?id=YsOyCpMUYD > - **Code:** https://github.com/uw-syfi/TeleRAG > [!abstract] 概要 > 検索拡張生成(RAG)は外部データソースと LLM を統合して事実の正確性とドメインカバレッジを向上させる技術である。現行の RAG パイプラインは大規模データストアに依存しており、GPU メモリが限られた環境で高スループットかつ低レイテンシを両立することが難しいというシステム上の課題を抱える。本論文では、レイテンシを削減し最小限の GPU メモリ要件でスループットを向上させる効率的な推論システム TeleRAG を提案する。中核技術は「ルックアヘッドリトリーバル」と呼ばれるプリフェッチ機構であり、リトリーバルに必要なデータを予測して LLM 生成と並行に CPU から GPU へ転送する。加えて、バッチ処理とマルチ GPU 推論を支えるプリフェッチスケジューラおよびキャッシュアウェアスケジューラを備える。評価では、シングルクエリで最大 1.98 倍のエンドツーエンドレイテンシ削減、バッチ処理で平均 1.83 倍のスループット向上を達成した。 ## 問題設定：RAG パイプラインのメモリとレイテンシの壁 - 現代の RAG パイプラインは 4 段階からなる：(1) 事前リトリーバル生成（クエリ変換）、(2) リトリーバル（ベクトルインデックス検索）、(3) 事後リトリーバル生成（応答生成）、(4) 判定 - データストアサイズは数十 GB から数 TB に達し、精度はデータストア規模と正の相関を持つ。wiki_dpr データセットで構築したインデックスは 61 GB - GPU メモリにインデックス全体を載せる方式（Case A）はレイテンシが低いがコストが高い。例：61 GB インデックス + 16 GB の Llama-3-8B = 77 GB で RTX4090（24 GB）には非現実的 - CPU オフロード方式（Case B）はメモリコストが低いがレイテンシが高い。CPU リトリーバルは全レイテンシの 41.1%（バッチサイズ 1）から 60.5%（バッチサイズ 4）を占める - GPU リトリーバルは CPU の 5.96 倍（バッチ 1）から 3.87 倍（バッチ 4）高速だが、メモリ消費が大きい ## 核心的洞察：ステージ間のクエリ相関 - RAG パイプラインにおいて、事前リトリーバル段階への入力クエリ $q_{\text{in}}$ とリトリーバル段階へ渡される出力クエリ $q_{\text{out}}$ の間には大きな意味的重複がある - 両クエリが同じ情報ニーズを表すため、対応する IVF クラスタの重複率は平均約 82% に達する - 3 つの QA データセット（NQ, HotpotQA, TriviaQA）と 6 つの RAG パイプラインで検証。nprobe=256 のとき、最低でも SubQ の 61.6% から S-RAG の 100% までの被覆率を確認（Table 1） | データセット | HyDE | SubQ | Iter | IRG | FLARE | S-RAG | |---|---|---|---|---|---|---| | NQ | 73.1% | 63.2% | 91.5% | 83.8% | 79.1% | 100.0% | | HotpotQA | 75.3% | 62.5% | 89.6% | 89.4% | 80.2% | 100.0% | | TriviaQA | 73.1% | 61.6% | 86.2% | 86.1% | 81.7% | 100.0% | ## 提案手法：ルックアヘッドリトリーバル - LLM による事前リトリーバル生成中に、入力クエリ $q_{\text{in}}$ の埋め込みを用いて必要な IVF クラスタを予測し、CPU から GPU へ非同期に DMA 転送する - リトリーバル時にハイブリッド検索を実行： 1. **予測・プリフェッチ**：LLM 生成中に $q_{\text{in}}$ との距離に基づき IVF クラスタを特定し、GPU DMA で非同期プリフェッチ 2. **GPU 類似度検索**：$q_{\text{out}}$ が利用可能になった時点で、プリフェッチ済みクラスタ（$C_{\text{overlap}}$）を GPU 上で高速検索 3. **CPU 類似度検索**：プリフェッチされなかったクラスタ（$C_{\text{miss}}$）は CPU で並行検索 4. **マージ**：GPU と CPU の検索結果を GPU 上でマージし、LLM に供給 - プリフェッチ量はクラスタ数でなくバイト予算 $b_p$ で制御。転送時間の予測可能性を確保するため固定バイト予算を採用 - 最適なプリフェッチ量は LLM の事前リトリーバル生成時間 $t_{\text{LLM}}$ に等しくなるまで転送する設計。プリフェッチ量 = $B_{\text{link}} \times \bar{t}_{\text{LLM}}$（$B_{\text{link}}$ はバス帯域幅） - GPU ソーティング：$C_{\text{miss}}$ のスカラー距離値のみを GPU に転送して $C_{\text{overlap}}$ と結合ソートを行い、軽量に最終結果を得る ## バッチ処理とマルチ GPU 対応 ### プリフェッチスケジューラ - バッチ処理ではクエリごとに異なる IVF クラスタが必要となり、プリフェッチ予算を共有するためヒット率が低下する - 解決策：意味的類似度の高いクエリをマイクロバッチにグループ化し、プリフェッチ対象クラスタの重複を最大化 - L2 距離に基づく貪欲探索でクエリをグループ化。バッチサイズ 256 以下では探索レイテンシが 0.1 秒未満であり、エンドツーエンドレイテンシへの影響は最小限 ### キャッシュアウェアスケジューラ - 頻繁にアクセスされる IVF クラスタを GPU メモリに保持するダイナミックキャッシングを実装 - マルチ GPU 環境では、各 GPU のキャッシュ内容とマイクロバッチのクラスタ要求の重複度に基づき、キャッシュローカリティを最大化するようマイクロバッチを GPU に割り当て - プリフェッチスケジューラのオーバーヘッドは約 37 ms、キャッシュアウェアスケジューラは約 180 ms（いずれも 4 GPU, バッチ 128 の構成で測定） ### マルチ GPU 設計 - グローバルバッチをマイクロバッチに分割し、各 GPU が独立に RAG パイプライン全体を処理するデータ並列方式を採用 ## 実験設定 - **データストア**: wiki_dpr（21 億トークン）に基づく 61 GB の IVF インデックス、4096 クラスタ、Contriever による 768 次元の埋め込み - **LLM**: Llama-3.2-3B、Llama-3-8B、Mistral-Small-22B - **RAG パイプライン**: HyDE, SubQuestion, Iterative, Iter-RetGen, FLARE, Self-RAG の 6 種 - **データセット**: NQ, HotpotQA, TriviaQA（各 1024 クエリをランダムサンプリング） - **ハードウェア**: | 構成 | Desktop | Server1 | Server2 | |---|---|---|---| | CPU | Threadripper 5975 | EPYC 9554 | EPYC 9534 | | CPU メモリ | 512 GB | 1.5 TB | 1.5 TB | | GPU | RTX4090 (24 GB) | H100 (80 GB) | 8 x H200 (140 GB) | | バス帯域幅 | PCIe 4 (32 GB/s) | PCIe 5 (64 GB/s) | PCIe 5 (64 GB/s) | - **ベースライン**: Faiss (CPU) + SGLang による CPU オフロード構成 - **パラメータ**: nprobe=256（$4\sqrt{4096}$）、top-k=3 ## 実験結果 ### シングルクエリレイテンシ（RTX4090） - Llama-3.2-3B 使用時、全パイプラインで CPU ベースラインを一貫して上回り、NQ で平均 1.55 倍、HotpotQA で 1.54 倍、TriviaQA で 1.49 倍の高速化 - Llama-3-8B 使用時も平均約 1.3 倍の高速化を維持。ピーク性能は Iter-RetGen の HotpotQA で 2.11 倍 - Llama-3-8B（16 GB）+ エンベディングモデル（1 GB）+ KV キャッシュ等を差し引いた残り 3.75 GB の GPU メモリのみでリトリーバルを加速 ### バッチスループット（H100） - Llama-3-8B, NQ データセットにおいて、バッチサイズ 1 で平均 1.32 倍、バッチサイズ 8 で平均 1.98 倍のスループット向上 - Mistral-Small-22B でもバッチサイズ 8 で平均 1.49 倍の向上 - バッチサイズ増大に伴いスループット向上率が拡大（CPU ベースラインはバッチサイズに対して線形にレイテンシが増加するのに対し、TeleRAG は GPU 並列性を活用） ### マルチ GPU スループット（H200 x 8） - グローバルバッチ 128、マイクロバッチ 4 の設定で、NQ における平均スループット倍率は 1 GPU 比で 2 GPU: 2.0 倍、4 GPU: 3.8 倍、8 GPU: 6.5 倍 - サブリニアスケーリングの主因はマイクロバッチ間の実行時間分散によるロードインバランス ### プリフェッチ予算とクラスタヒット率 - プリフェッチ予算が大きいパイプライン（HyDE: H100 で 10 GB、ヒット率 93.2%）ではヒット率 50% 超を安定的に達成 - 予算 2 GB 未満のケース（IRG: RTX4090 で 2.5 GB、ヒット率 45.1%）でも GPU ソーティングと CPU 計算負荷削減の複合効果により 1.2-1.6 倍の高速化を実現 ### キャッシュのアブレーション - キャッシュによるスループット向上は GPU 数に応じて増大：1 GPU で 2%、2 GPU で 12%、4 GPU で 21%、8 GPU で 18% - 単一 GPU では限られたキャッシュ容量で多様なクエリに対応しきれないが、GPU 数が増えるとキャッシュアウェアスケジューリングが有効に機能 ## 関連研究との差異 - **KV キャッシュ最適化**（RagCache, TurboRAG, CacheBlend 等）：プリフィルレイテンシの削減が主目的であり、リトリーバルレイテンシには対処しない。TeleRAG はこれらと相補的 - **投機的リトリーバル**（Zhang et al., 2024b）：トークン単位の投機的デコーディングと組み合わせた手法であり、モジュラー RAG パイプラインには適用不可 - **HedraRAG, RAGO, Hermes**：同時期の並行研究。HedraRAG はリクエスト内類似性を活用、RAGO はスキーマベースの自動性能解析、Hermes は複数 CPU ノードへの分散検索をそれぞれ提案 - **EdgeRAG, LEANN**：メモリ削減に特化するがリトリーバル加速は限定的。TeleRAG は IVF インデックスで加速とメモリ削減の両方を達成 ## 議論：適用範囲の拡張 - **IVF 以外のインデックス**: LSH（ハッシュバケットを IVF クラスタと見做すことで自然に拡張可能）、HNSW（グラフ走査ベースのため直接適用は困難だが、FAISS の IVF-HNSW ハイブリッド構造を通じた適用が可能） - **新興ハードウェア**: NVIDIA Grace 等の高帯域システムではプリフェッチ量を増やしてヒット率を向上可能。Apple M シリーズ等の統合メモリ環境では TeleRAG の有用性は低下するが、現状のメモリ容量制約下では依然有効 - **メモリ割り当てトレードオフ**: LLM と KV キャッシュ用メモリを優先確保し、残りをリトリーバルに割り当てる設計。リトリーバルが LLM のバッチサイズやシーケンス長を制限しないことを保証 ## 結論 - TeleRAG は、ルックアヘッドリトリーバルによる CPU-GPU データ転送の隠蔽、プリフェッチスケジューラによるバッチ最適化、キャッシュアウェアスケジューラによるマルチ GPU 効率化を統合した RAG 推論システムである - RTX4090 上でシングルクエリ平均 1.53 倍（Llama-3.2-3B）の高速化、H100 上でバッチサイズ 8 時に平均 1.98 倍のスループット向上、8 GPU（H200）で 1 GPU 比 6.5 倍のスケーリングを達成した - ACM Artifacts Available / Functional / Results Reproduced の 3 バッジを取得