AI-ügynökök lekérdezése akár 8-szeres sebességgel egyedi GPU-algoritmussal

Anubhab Banerjee egy blogbejegyzésében bemutatott egyedi CUDA-kernel megoldást, amely drasztikusan csökkenti az AI-ügynökök lekérdezési folyamatainak (RAG - Retrieval-Augmented Generation) késleltetését. A fejlesztő szerint a hagyományos architektúrákban a GPU és a CPU közötti adatforgalom jelenti a fő szűk keresztmetszetet, nem maga a számítás.

A „round-trip tax” és az adatforgalmi probléma

A probléma lényege, hogy az ügynökök minden egyes kontextuskereséskor lekérdezést küldenek a GPU-ról a Python felé. A CPU pontozza az adatokat, kiválasztja a legjobb K elemet, majd visszaküldi az eredményt a GPU-ra. Ez az oda-vissza út, különösen nagy adatbázisok esetén, jelentős késleltetést okoz. Banerjee szerint ez a „round-trip tax” a felelős az ügynökök lassulásáért.

Kapcsolódó: GPU-kernelek Pythonból

A gyorsítás technikája és eredménye

A megoldás kulcsa, hogy az adatbázist egyszer feltöltik a GPU memóriájába (VRAM), és a hasonlósági keresést, a Top-K kiválasztást és az eredmények összevonását teljes egészében a GPU-n végzik. Csak az apró lekérdezési vektor és a K eredmény utazik a PCIe-buszon. Ez a megközelítés még egy 7 éves GTX 1080-as GPU-n is akár 8,57-szeres gyorsulást eredményezett az optimalizált CPU-s alapkonfigurációkhoz képest.

Kapcsolódó: CUDA Tile C++ támogatás

A tesztek során K=8 értéknél minden 15 vizsgált konfigurációnál (N ∈ {10k, 50k, 100k, 500k, 1M}, D ∈ {384, 768, 1024}) sebességnövekedést mértek, 2,43× és 8,57× között. K=32 értéknél 13 konfigurációnál volt gyorsabb a GPU-s megoldás, elérve a 7,76× sebességnövekedést. K=100 értéknél azonban a CPU-s alapkonfiguráció bizonyult gyorsabbnak 15-ből 14 esetben, ami a szerző szerint a kernel egyszerűsége miatt történt.

Kapcsolódó: JetPack 7.2 AI-telepítés

A fejlesztő hangsúlyozza, hogy a sebességnövekedés nem a „varázskernelnek” köszönhető, hanem annak, hogy elkerülték a szükségtelen adatmozgatást a GPU és a fő memória között. Az új megoldás egy kis C++ orchestrator API-ként érhető el, amely lehetővé teszi az adatbázis feltöltését és a lekérdezések futtatását.

Kapcsolódó: GPU-flottafelügyelet

Banerjee, aki korábban 5G/6G RAN mérnökként dolgozott, a beam selection folyamatokat hasonlította a RAG Top-K kereséshez. A GPU-n tartott corpus és a teljes folyamat GPU-n futtatása drasztikusan csökkenti a késleltetést, ami kulcsfontosságú lehet a komplex, több lépésből álló AI-ügynökök teljesítménye szempontjából.

Kapcsolódó: Rust CUDA fordítás

A CUDA-TopK-Retrieval megoldás a „Production-Grade Agentic Inference” sorozat harmadik része, amelynek célja a redundáns számítási és kommunikációs terhek eltávolítása az AI-folyamat-okból. A CUDA-TopK-Retrieval kódtár jelenleg elérhető a GitHubon.