NVIDIA bemutat egy teljes, éles környezetre tervezett útvonalat Hierarchical Sequential Transduction Unit (HSTU) alapú generatív ajánlórendszerek üzemeltetésére. A bemutatott munkafolyamat kombinálja a HSTU modellt, a PyTorch Ahead‑of‑Time Inductor (AOTI) használatával történő exportálást és fordítást, GPU‑támogatású (NV) embedding cache‑t, FlexKV alapú KV cache szolgáltatást, natív C++ validálást és a Dynamo‑Triton üzemeltetési réteget. A cél: alacsony késleltetésű inferencia olyan hosszú felhasználói történetek és nagy kategóriájú embedding táblák mellett, amelyek jellemzőek a modern személyre szabásra.
Miért HSTU a generatív ajánláshoz?
A generatív ajánlórendszerek (GR) ajánlást nem különálló lekérdezési, rangsorolási és előrejelzési lépések sorozataként kezelik, hanem a felhasználói események sorrendi modellezéseként. A HSTU erre a feladatra készült: a modell bemenete tokenek sorozata, amely tartalmazza a kontextust, a felhasználói interakciókat, az elemeket és esetleges akciótokeneket. Az előfeldolgozás embeddingeket kér le, beilleszti az item és action embeddingeket, hozzáadja a kontextuális tokeneket és pozíciós kódolást, majd HSTU blokkok dolgozzák fel a sorozatot, végül multitask rangsorolási kimeneteket ad a predikciós fej.
Ez a felépítés jól kezeli a recency‑t, a sorrendet és az ismétlődő interakciós mintákat — ugyanakkor inferencia során költséges lehet, ha minden lekérésnél hosszú történeteket kell újraszámolni.
Miért nehéz nagy, sorrendi ajánlórendszereket szolgáltatni?
A szolgáltatásrétegnek kezelnie kell a különböző hosszúságú (jagged) sorozatokat, hatalmas kategórikus embedding állapotot és azt a gyakorlatot, hogy ugyanaz a felhasználó gyakran visszatér csak néhány új tokennel. Ha minden kérésnél újraépítjük a teljes key‑value állapotot, az fölösleges műveleteket és késleltetést eredményez.
A KV cache szerepe, hogy tárolja a korábban számított kulcs‑érték blokkokat, így a modell elkerülheti a cache‑elt prefix újraszámítását — különösen hasznos, ha a felhasználó hosszú története nagyrészt változatlan marad.
A Dynamo‑Triton HSTU inferencia munkafolyamata
A bemutatott inferenciaútvonal tartalmaz egy KVCacheManager komponenst, amely GPU memóriát és host‑oldali tárolót használ a paged KV táblázat megvalósításához. A GPU cache lookup/alloc/append/evict műveleteket támogat, LRU‑szerű elv alapján kiürítéssel, ha helyhiány lép fel. Host‑oldali réteg további tárolást nyújt, a FlexKV pedig backend runtime‑ként kezeli a cache‑elt adatokat.
A HSTU attention kernel képes a paged cache‑ből fogyasztani a KV adatokat; az exportált inferenciaút tartalmaz cache‑tudatos custom opokat (lookup, allocation, onboarding, append, offload), így megőrizhető a sorozati szemantika miközben csökkentjük a redundáns számításokat.
PyTorch AOTI: natív, előre lefordított inferencia
A PyTorch AOTI (Ahead‑of‑Time Inductor) folyamat a PyTorch modellt torch.export és AOTI segítségével exportálja és előre lefordítja. Az eredmény egy .pt2 csomag plusz metaadatok és embedding táblák, amely betölthető natív C++ runtime alatt, így csökken a Python futtatási overhead és deployment‑barát artifact jön létre.
Az NVIDIA példa kombinálja a DynamicEmb inference embedding táblákat az NV Embedding Cache‑cel: csak a népszerű embeddingek kerülnek GPU memóriába, a teljes tábla CPU memóriában marad. Az export út pedig metaadatokat és embedding adatokat ír mellé, elkerülve a felesleges többszörös embedding másolatokat.
Az exportált csomagokat többszörösen validálják: Python export és replay, valamint natív C++ replay az érvényesség és teljesítmény ellenőrzésére. Ugyanezen AOTI csomagot használja a Dynamo‑Triton is, ami segít, hogy a fejlesztői validálás és a production serving közelebb legyen egymáshoz.
Dynamo‑Triton telepítési út
A Dynamo‑Triton szolgáltatja a production szintű kéréskezelést, model repository‑t, backend integrációt, mérőszámokat és üzemeltetési struktúrát. Az AOTI deployment a Dynamo‑Triton PyTorch backendjét használja platform: "torch_aoti" beállítással, így a Dynamo‑Triton képes betölteni és kiszolgálni az előre lefordított PyTorch csomagot.
A teljes munkafolyamat öt lépése:
- A szükséges custom operátorok és runtime könyvtárak buildelése
- HSTU modell exportálása PyTorch AOTI‑val
- FlexKV alapú KV‑cache szolgáltatás elindítása
- Az exportált artefaktok natív C++ replay‑jel történő validálása
- Az exportált KV‑cache AOTI modell kiszolgálása Dynamo‑Tritonon
Az NV Embedding Cache csökkenti a GPU memóriaigényt, a FlexKV csökkenti az ismételt számítást, a Dynamo‑Triton pedig a gyártási környezet kívánt képességeit biztosítja.
Benchmarkok és mérési eredmények
A demonstrációs benchmarkok a NVIDIA/recsys‑examples repó KuaiRand‑1K konfigurációját használják egyetlen NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU‑n. A mérések paraméterei:
- Modellek: három‑rétegű és nyolc‑rétegű HSTU variánsok
- Rejtett méret (hidden size): 512
- Attention headek: 4
- Modell súlyok: BF16
- KV cache: BF16
- Maximális történet hossz (history): 8,192 token (4,096 item+action párok)
- Maximális jelölt (candidate) sorozat: 100 token
- Kontextuális feature‑ok: 6
- Exportált effektív sorozathossz (aligned): 8,320 token
A benchmark protokoll latency‑t számol logikai kérésre vetítve, kizárva az adatbetöltést, validálást, rebatchinget, szerver indítást, warmup időszakot és post‑warmup szünetet a mért időből.
Főbb eredmények:
- Dynamo‑Triton batch size 2 mellett a PyTorch AOTI backend már KV‑cache nélkül is jobb késleltetést ad a Dynamo‑Triton Python backendhez képest; 20 GB GPU KV cache esetén a javulás lényegesen nagyobb.
- Batch size 8 mellett, 100% GPU KV‑cache találati aránnyal, a Dynamo‑Triton + PyTorch AOTI + FlexKV elérte a legjobb esetben:
- 4.47× gyorsulás a háromrétegű HSTU modellnél (összehasonlítva ugyanazzal az AOTI konfigurációval KV cache nélkül)
- 5.93× gyorsulás a nyolcrétegű modellnél
- Konkrét latenciaértékek batch size 8 mellett: háromrétegű modell ~0.423 ms / logikai kérés (GPU KV‑cache hit), nyolcrétegű modell ~0.678 ms / logikai kérés (GPU KV‑cache hit).
Ezek az eredmények azt mutatják, hogy a fordított, cache‑tudatos üzemeltetés különösen hatékony mélyebb modelleknél és nagyobb logikai batch‑méreteknél, mert a több rétegen átívelő ismételt számítást sikerül elkerülni.
Milyen előnyöket ad az HSTU gyorsítása?
Ajánlórendszerek szoros késleltetési korlátok alatt működnek: a ranking késleltetése befolyásolhatja az oldalbetöltést, a feed válaszkészségét és a hirdetéskiszolgálást. Az egyre személyre szabottabb, sorrendérzékeny modellek nagyobb inferencia‑igényt támasztanak.
A bemutatott stack (HSTU + PyTorch AOTI + FlexKV + NV Embedding Cache + Dynamo‑Triton) egyszerre csökkenti a futtatási overheadet, csökkenti a GPU memóriaigényt és minimalizálja a redundáns attention számításokat, különösen amikor a felhasználói történelem prefixe változatlan marad több lekérés között.
Hogyan kezdheted el?
A teljes munkafolyamat reprodukálható és kiterjeszthető a NVIDIA/recsys‑examples GitHub repóból. A repo tartalmaz AOTI inference útmutatót, HSTU áttekintést, telepítési leírásokat, KuaiRand‑1K adatelőkészítést, checkpoint tréninget, export‑ és C++ replay példákat, valamint Dynamo‑Triton image‑csomagolási útmutatót.
Köszönetnyilvánítás
A publikáció több NVIDIA csapat keresztfunkcionális együttműködésének eredménye. Köszönet jár J, Runchu Zhao, Yulu Liu, Lin Hu, Zhuofan Li, Jacob Subag és Tomer Bar‑On közreműködéséért.
