Atskirti „prefill“ ir „decode“

Siekiant generuoti išvesties žetonus iš įvesties užklausos, LLM inferencija yra padalinta į du etagus: „prefill“ (pirminio užpildymo) ir „decode“ (dekodavimo). „Prefill“ apdoroja įvesties žetonus, užpildydamas KV talpyklas, prieš pereinant į „decode“ etapą, kuris generuoja žetonus po vieną. Nors vienas dekodavimo žingsnis paprastai trunka dešimtis milisekundžių, „prefill“ užtrunka gerokai ilgiau. Vykdant juos ant tų pačių įrenginių, „prefill“ maišymas su „decode“ pablogina dekodavimo našumą. Šiame straipsnyje nagrinėjamas nusistovėjęs sprendimas – atskirtas „prefill“ ir „decode“, vykdant juos ant atskirų įrenginių, kad būtų maksimaliai padidintas tiek „prefill“ pralaidumas, tiek dekodavimo delsa.

„Prefill“ ir „Decode“ našumas

Tipiniame LLM tarnybinės stoties variklyje paketų planuoklis atrenka užklausas apdorojimui kiekviename modelio vykdymo žingsnyje. Vykdant ant vieno įrenginio arba mazgo, tiek „prefill“, tiek „decode“ užklausos sugrupuojamos į paketus. Dėmesio mechanizmo („attention“) sąnaudos, kurios didėja kartu su sekos ilgiu, auga tiek „prefill“, tiek „decode“ atveju proporcingai įrašų ilgiam KV talpykloje (kv_len). „Decode“ užklausos paprastai perduoda vieną žetoną (qo_len=1) su minimaliomis sąnaudomis per kitus sluoksnius, kurie nepriklausomai veikia su sekos žetonais. „Prefill“ užklausos perduoda tūkstančius ar dešimtis tūkstančių žetonų su didelėmis sąnaudomis per tankiuosius sluoksnius (didelis qo_len).

Tiesioginio pralaidumo delsa labiau priklauso nuo nepriklausomų žetonų, perduotų per tankiuosius sluoksnius (qo_len), skaičiaus nei nuo žetonų, gautų iš KV talpyklos dėmesio mechanizmo metu (kv_len), skaičiaus. Dėmesio mechanizmas gali lygiagretinti tiek pagal užklausų skaičių, tiek pagal kv_len, proporcingą sekų ilgiams, taip pasiekdamas gerą išnaudojimą. „Prefill“ yra apribotas skaičiavimo resursų: esant dideliam qo_len, GEMM branduoliai gali paskankinti pakankamai blokų per M matmenį, kad visiškai išnaudotų šiuolaikinių GPU skaičiavimo galimybes. „Decode“ yra apribotas atminties: dėl paprastai mažų paketų dydžių įvestis skaičius pagal M dažniausiai yra mažas, pakankamas tik vienam blokui. Nors „Split-K GEMM“ branduoliai gali pagerinti SM išnaudojimą mažiems žetonų paketų dydžiams, talpyklos ir matricų daugybos blokai dažniausiai lieka išnaudojami nepakankamai.

Padidinti

Kai jie yra sumaišyti, paketai, kuriuose yra „prefill“ užklausų, patiria didesnę delsą atliekant tiesioginį perdavimą, o tai neigiamai veikia viso egzemplioriaus dekodavimo pralaidumą. Nors „prefill“ užklausų maišymas su „decode“ užklausomis arba dalinio „prefill“ naudojimas gali šiek tiek pagerinti dekodavimo našumą, sunku išlaikyti pakankamą „prefill“ pralaidumą, kad būtų apdorota pakankamai užklausų egzemplioriuje ir maksimaliai padidintas dekodavimo pralaidumas. Didelių modelių atveju, esant tipiniams išvesties ilgiams, siekiant išlaikyti didelį paketo dydį dekodavimui, „prefill“ turi būti atliekamas pakankamai dažnai, kad tai žymiai pablogintų vidutinę delsą ir sukeltų strigimą išvestyje.

Padidinti

Padidinti

Šias problemas galima išspręsti naudojant atskirą mazgų rinkinį „prefill“ ir „decode“ atlikimui. Susiejus „prefiller“ mazgą su keliais „decoder“ mazgais, galima suplanuoti pakankamai užklausų „prefill“ etapui, kad būtų maksimaliai padidintas pralaidumas ir išlaikytas pakankamai didelis vienu metu vykstančių užklausų skaičius „decoder“ mazguose, taip pat maksimaliai padidinant dekodavimo pralaidumą. „Prefiller“ mazgai užpildo KV talpyklas, kurios vėliau perduodamos „decoder“ mazgams. Kadangi dekoderiai nebeprivalo sustoti dėl „prefill“, delsa tampa kur kas labiau nuspėjama, nes bendras augančio aktyvių užklausų kv_len poveikis yra kur kas mažiau pastebimas. Kaina sumokama pailgėjus laiko iki pirmojo žetono (TTFT) rodikliui, nes KV talpyklų perkėlimas tinkle gali užtrukti nuo dešimčių iki šimtų milisekundžių.

KV Messenger

„Perplexity“ mūsų atskirto „prefill“ ir „decode“ įgyvendinimas yra sukurtas aplink KV pranešimų siuntiklį („KV messenger“), kuris sąveikauja su LLM varikliu, kad orkestruotų KV talpyklos perkėlimus iš „prefiller“ mazgų į „decoder“ mazgus per tinklą. „Prefiller“ pusėje pranešimų siuntiklis priima užklausas iš „decoder“ mazgų, perduoda jas paketų planuokliui ir seka tiesioginio perdavimo vykdymą, kad išsiųstų KV talpyklas su kuo mažesne delsa. „Decoder“ pusėje, po to, kai paskirstomi neiškeliami puslapiai, pranešimų siuntiklis blokuoja užklausos planavimą dekodavimui, kol gaunamas pranešimas apie KV talpyklos ir dekoderio konteksto perkėlimų pabaigą.

Padidinti

Atskirtas „prefill“ reikalauja didelio pralaidumo ir mažos delsos ryšių, todėl mūsų sprendimas pritaikytas RDMA, palaikantis tiek EFA, tiek „ConnectX“ tinklo sąsajos valdiklius (NIC). KV Messenger yra sukurtas naudojant libfabric, pasitelkiant mūsų fabric-lib apvalkalus, kad būtų suteiktos aukštesnio lygio mažos delsos abstrakcijos virš nuotolinės tiesioginės atminties prieigos (RDMA) primitivų, įgyvendinant efektyvų puslapių ir metaduomenų perkėlimą kartu su mažos delsos signalizavimu. Fone fabric-lib koordinuoja GPU ir jo tiesiogiai prijungtus NIC, kad nukopijuotų duomenis iš „prefiller“ mazgo į „decoder“ mazgą.

Gavęs užklausą, „prefiller“ mazgas paskirsto atitinkamą šaltinio KV puslapių rinkinį ir suplanuoja „prefill“ užklausą naudodamas savo vietinį variklį. Siekiant sumažinti delsą, perkėlimai nelaukia tiesioginio perdavimo: vietoj to KV puslapių kopijavimas pradedamas, kai tik modelis baigia pridėti KV talpyklos įrašus į KV talpyklą atskiriems sluoksniams. Kadangi „prefill“ užklausos gali būti skaidomos dalimis, paketų planuoklis praneša KV pranešimų siuntikliui apie šiuo metu suplanuotas dalis prieš vykdymą. Kad būtų galima palaikyti CUDA grafikus ir kartu sekti sluoksnius, pranešimų siuntiklis palaiko dedikuotą giją, kuri tikrina skaitiklį, padidinamą po dėmesio mechanizmo išvesties projekcijos. Skaitiklis palaikomas tik pagrindiniame mazge šarduotoje aplinkoje: nors KV talpyklos įrašai yra galiojantys po pridėjimo ir prieš dėmesio mechanizmą, išvesties projekcija yra sumažinama per rangus, juos sinchronizuojant. Kai pastebimas skaitiklio pasikeitimas, apie tai pranešama pranešimų siuntikliui ir jis iškviečia fabric-lib, kad pradėtų sluoksnio perkėlimą.

Padidinti

Baigus paskutinės dalies perkėlimą, taip pat nukopijuojami visi papildomi metaduomenys: spekuliatyvusis dekodavimas arba MTP reikalauja, kad logitai ir paslėptos būsenos būtų perkeltos į dekoderį. Šios kopijos taip pat atliekamos per RDMA į ir iš iš anksto paskirstytų buferiui skirtų srauto vietų.

Baigus visus laukiančius paskutinės dalies perkėlimus, „prefiller“ mazgas atšaukia KV puslapių paskirstymą ir užbaigia užklausą. „Decoder“ mazgui apie tai aiškiai nepranešama: vietoj to jis naudoja tiesioginius skaitiklius, kad sektų atliktų operacijų skaičių. RDMA operacijų skaičius „prefiller“ pusėje yra proporcingas perkeltų puslapių skaičiui. Baigus žinomą puslapių ir konteksto kopijų skaičių, fabric-lib iškviečia KV pranešimų siuntiklį, kad nurodytų, jog užklausos yra paruoštos dekodavimui. Pranešimų siuntiklis atšaukia bet kokio konteksto paskirstymą ir perduoda užklausą LLM varikliui.

Šarduoti KV talpyklos perkėlimai

Jei „prefiller“ ir „decoder“ remiasi tenzorių lygiagretumu (TP) ir šarduoja arba identiškai replikuoja KV talpyklas, vienas perkėlimo variklis koordinuoja kelis įrenginius, kad išsiųstų ir gautų visų replikų puslapius. Siekiant galimybės naudoti vieną pranešimų siuntiklį ir perkėlimo variklį nepaisant to, kad modelio vykdytojas yra replikuojamas per kelis įrenginius ir procesus, cuMem ir cuMemImportFromShareableHandle yra naudojami įrenginio atminčiai, palaikančiai KV talpyklas, paskirstyti ir susieti ją su pagrindiniu procesu. Perkėlimo variklis patikrina mazgo topologiją, kad surastų NIC ir procesorius, esančius artimiausiame NUMA mazge, kuriuos naudoti kiekvieno KV talpyklos pjūvio perkėlimui.

Jei šaltinis ir paskirties vieta šarduoja identiškai, perkėlimai yra paprasti, nes egzistuoja vienas su vienu atitiktis tarp šaltinio ir paskirties vietos įrenginių bei puslapių. Tokiu atveju šardavimas netiesiogiai padeda perkėlimo delsai: naudojant daugiau GPU, galima pasitelkti daugiau susietų NIC, pasiekiant artimesnį pilno pralaidumo išnaudojimą. Tačiau, jei yra neatitikimas, perkėlimo variklis turi padalinti arba atkurti puslapius priklausomai nuo santykio tarp šaltinio ir paskirties vietos pjūvių.

Padidinti

Jei „prefiller“ padalija KV talpyklą per daugiau įrenginių, pilni puslapiai atkuriami dekoderyje, išsiunčiant atitinkamas puseles iš „prefiller“ įrenginių. Jei dekoderis turi daugiau šardų, jis gauna puslapius iš kelių šaltinių. Dekoderis turi žinoti „prefiller“ šardavimo schemą, kad galėtų apskaičiuoti RDMA įrašų, kuriuos jis tikisi gauti, skaičių. Jei taikoma replikacija, „prefiller“ sugrupuoja įrenginius į replikų rinkinius, kurie patys replikuoja visą KV talpyklą. Paskirties replikų rinkiniams atsitiktine tvarka priskiriamas vienas iš šaltinio rinkinių, siekiant panaudoti visus turimus įrenginius RDMA įrašams inicijuoti.

Padidinti

Šarduoti perkėlimai reikalauja nedidelio KV talpyklų pakoregavimo. Pagal nutylėjimą „FlashInfer“ remiasi NHD išdėstymu, kuris išdėsto žetonus puslapyje galvutėse. Kadangi talpyklos greičiausiai yra šarduojamos pagal dėmesio galvučių skaičių, tai sukuria trūkumą galvutės viduje. RDMA perkėlimai netiesiogiai nepalaiko įrašų su žingsniu, todėl perkėlimui atlikti reikia vienos operacijos galvutei. Vietoj to, siekiant sumažinti sąveikos su libfabric skaičių, mes organizuojame KV talpyklas naudodami HND išdėstymą, kuris pastato galvutės matmenį prieš žetonų skaičių. Tai užtikrina tęstinumą, leidžiantį nukopijuoti puslapį vienu įrašu.

Spekuliatyvusis dekodavimas

Spekuliatyvusis dekodavimas reikalauja nedidelių pakeitimų atskirtame „prefill-decode“. Mūsų įgyvendinime „prefiller“ mazgams neleidžiama atlikti žetonų atrankos („sample“). Kadangi „Perplexity“ „Sonar“ modeliai palaiko struktūrizuotą išvestį, nenorime patirti schemos procesoriaus įgyvendinimo sinchronizavimo tarp „prefiller“ ir „decoder“ sudėtingumo. MTP ir spekuliatyviojo dekodavimo mechanizmuose juodraštinio modelio paruošimas iki paskutinio žetono apima žetonų atranką iš tikslinio modelio.

Padidinti

Norint išspręsti šias problemas, „prefill“ neapima paskutinio įvesties sekos žetono. Vietoj to perkeliamos paslėptos būsenos arba logitai iš „prefill“, buvusio prieš paskutinį žetoną, ir kitame žingsnyje dekoderyje jis traktuojamas kaip dekodavimo žetonas. Nors tai šiek tiek padidina delsą, kadangi po „prefill“ reikia atlikti pilną dekodavimo žingsnį pirmajam žetonui išleisti, įgyvendinimo sudėtingumas yra smarkiai sumažinamas.

Atskirti diegimai

Diegėme arba eksperimentavome su keliomis atskirtomis konfigūracijomis su įvairiais modeliais, kad palaikytume gamybinį srautą arba vidinius vertinimo darbo krūvius. Remdamiesi modelių dydžiu ir dėmesio mechanizmu, pasirinkome tinkamas šardavimo schemas „prefiller“ ir „decoder“ mazgams, kad optimaliai išnaudotume GPU.

DeepSeek-R1

Naudodami „DeepSeek“, svarstėme tiek tenzorių lygiagretumo (TP), tiek duomenų lygiagretumo (DP) diegimus. Kaip aptarta ankstesniuose tinklaraščio įrašuose, TP diegimai užtikrina geresnę delsą mažesnio pralaidumo sąskaita, reikalaujant daugiau GPU dideliam srautui aptarnauti. DP diegimai skalbyjasi daug geriau didėjant apkrovai, tačiau jų didžiausias pralaidumas yra mažesnis dėl įrenginių ar mazgų tarpusavio komunikacijos sąnaudų.

„DeepSeek“ remiasi daugiagalviu latentiniu dėmesiu („Multi-Head Latent Attention“), suspausdamas KV talpyklas. Kadangi visos KV galvutės yra suspaustos į vieną latentinį vektorių, TP negali šarduoti KV talpyklų, nes turi replikuoti latentinius vektorius visuose rankuose. Šardavimas įvyksta po dekompresijos, nes kiekvienas rankas gali išgauti skirtingas galvutes iš to paties latentinio vaizdavimo. Vadinasi, visi KV talpyklos šardai yra identiški tiek „prefiller“, tiek „decoder“ šarduose.

Naudojant mazgo vidinį TP sąranką, tiek „prefiller“, tiek „decoder“ yra šarduojami identiškai. Perkėlimai siunčiami iš visų rankų, siekiant visiškai išnaudoti visus turimus NIC. Tačiau naudojant DP diegimą, kur TP ranko dydis yra mažesnis arba kiekvienas DP rankas priskirtas vienam GPU, bet kuris „prefiller“ įrenginys, turintis replikuotą KV talpyklos kopiją, gali ją išsiųsti. Norėdami subalansuoti užklausas per visus turimus NIC, atsitiktine tvarka pasirenkame GPU ir NIC KV talpyklai išsiųsti iš „prefiller“ į „decoder“.

Esant sumaišytam „prefill-decode“, mūsų R1 diegimas sunkiai palaikė nuosekliai didesnį nei 50 TPS greitį dėl dažnų „prefill“ trukdžių šimtų milisekundžių ribose. Priešingai, atskirus „prefill“, patyrėme apie 100 ms TTFT bausmę kiekvienai užklausai, tačiau vienas „prefiller“ mazgas galėjo išlaikyti nuoseklius paketų dydžius 3 „decoder“ mazguose, pasiekdamas didesnį nei 90 TPS pralaidumą, aptarnaudamas apie 1 QPS apkrovą vienam „decoder“ mazgui. Duomenų lygiagretumo diegimuose TPS buvo šiek tiek mažesnis – apie 50, tačiau egzemplioriai galėjo apdoroti 1 QPS apkrovą vienam rankui, su 8 rankais viename mazge.

Qwen3-Coder

Šis 480B modelis naudoja grupinės užklausos dėmesį (GQA), todėl dėmesio mechanizmas gali būti lengvai šarduojamas ir gauti naudos iš tenzorių lygiagretumo neaukojant atminties KV talpykloms. Todėl galėjome šarduoti modelį per 8 GPU tiek „prefill“, tiek „decode“ etapams, suporuodami apie 3 „decoder“ mazgus su vienu „prefiller“ mazgu. Kadangi dėmesio mechanizmas yra šarduotas, mes remiamės HND KV talpyklos išdėstymu „prefiller“ ir „decoder“ KV talpykloms šarduoti, suporuodami „prefiller“ rankus su „decoder“ rankais ir visiškai išnaudodami visus NIC pjūviams perkelti lygiagrečiai.

Citata

@online{perplexity2025DisaggregatedPrefillAndDecode,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Kevin Hu and Erin McKnight and Phi Hoang},
  title = {{Disaggregated prefill and decode}},
  date = {2025-08-01},
  url = {https://perplexity.coolkeypoint.com/hub/blog/disaggregated-prefill-and-decode}
}