Sadalīta sagatavošana un dekodēšana

Lai ģenerētu izvades tokenus no ievades uzvednes, LLM secināšana tiek sadalīta divos posmos: sagatavošanā un dekodēšanā. Sagatavošana apstrādā ievades tokenus, aizpildot KV kešatmiņas, pirms tiek pārietts uz dekodēšanas posmu, kas ģenerē tokenus pa vienam. Lai gan viens dekodēšanas solis parasti ilgst desmiti milisekunžu, sagatavošana prasa ievērojami ilgāku laiku. Ja tos palaiž uz tām pašām ierīcēm, sagatavošanas sajaukšana ar dekodēšanu samazina dekodēšanas veiktspēju. Šajā rakstā mēs aplūkojam izveidotu risinājumu – sadalīto sagatavošanu un dekodēšanu, izpildot tos uz atsevišķām ierīcēm, lai maksimāli palielinātu gan sagatavošanas caurlaidspēju, gan dekodēšanas latentumu.

Sagatavošanas un dekodēšanas veiktspēja

Tipiskā LLM apkalpošanas dzinējā pakešu plānotājs atlasa pieprasījumus apstrādei katrā modeļa izpildes solī. Darbojoties uz vienas ierīces vai mezgla, gan sagatavošanas, gan dekodēšanas pieprasījumi tiek grupēti kopā. Uzmanības izmaksas, kas palielinās secības garumā, pieaug gan sagatavošanai, gan dekodēšanai – proporcionāli KV kešatmiņas ierakstu garumam (kv_len). Dekodēšanas pieprasījumi parasti pārsūta vienu tokenu (qo_len=1), radot minimālas izmaksas citos slāņos, kas neatkarīgi darbojas ar secības tokeniem. Sagatavošanas pieprasījumi pārsūta tūkstošiem vai desmitiem tūkstošu tokenu par ievērojamām izmaksām blīvajos slāņos (liels qo_len).

Latences ilgums priekšējās izpildes laikā vairāk ir atkarīgs no neatkarīgo tokenu skaita, kas tiek nodots caur blīvajiem slāņiem (qo_len), nevis no tokenu skaita, kas tiek iegūts no KV kešatmiņas uzmanības mehānisma laikā (kv_len). Uzmanības mehānisms var paralelizēt gan pa pieprasījumu skaitu, gan pa kv_len, kas ir proporcionāls secību garumiem, tādējādi nodrošinot labu resursu izmantošanu. Priekšpiepilde ir skaitļošanas ierobežota: ja qo_len ir augsts, GEMM kodoli var piešķirt pietiekami daudz bloku gar M dimensiju, lai pilnībā izmantotu moderno GPU skaitļošanas iespējas. Dekodēšana ir atmiņas ierobežota: parasti zemā pakešu izmēra dēļ ievadu skaits gar M parasti ir mazs, pietiekams tikai vienam blokam. Lai gan Split-K GEMM kodoli var uzlabot SM izmantošanu pie zema tokenu pakešu izmēra, kešatmiņas un matricas reizināšanas vienības parasti paliek nepietiekami izmantotas.

Izvērst

Apvienojot tos, paketes, kas satur sagatavošanas pieprasījumus, rada lielāku latentumu priekšējā caurlaidē, negatīvi ietekmējot visas instances dekodēšanas caurlaidspēju. Lai gan sagatavošanas pieprasījumu sajaukšana ar dekodēšanas pieprasījumiem vai sadalītās sagatavošanas izmantošana var nedaudz uzlabot dekodēšanas veiktspēju, ir grūti uzturēt pietiekamu sagatavošanas caurlaidspēju, lai apstrādātu pietiekami daudz pieprasījumu uz instances, lai maksimāli palielinātu dekodēšanas caurlaidspēju. Lielu modeļu gadījumā ar tipiskiem izvades garumiem, lai saglabātu lielu paketes izmēru dekodēšanai, sagatavošana ir jāveic pietiekami bieži, kas būtiski pasliktina vidējo latentumu un rada saraustītību izvade.

Izvērst

Izvērst

Šīs problēmas var novērst, izmantojot atsevišķu mezglu kopu sagatavošanas un dekodēšanas veikšanai. Piesaistot sagatavošanas mezglu vairākiem dekodētāja mezgliem, var ieplānot pietiekami daudz pieprasījumu sagatavošanai, lai maksimāli palielinātu caurlaidspēju un saglabātu pietiekami lielu vienlaicīgu pieprasījumu skaitu dekodētāja mezglos, lai maksimāli palielinātu arī dekodēšanas caurlaidspēju. Sagatavošanas mezgli aizpilda KV kešatmiņas, kas pēc tam tiek pārsūtītas uz dekodētāja mezgliem. Tā kā dekodētājiem vairs nav jāpārtrauc darbs sagatavošanas dēļ, latentums kļūst daudz deterministiskāks, jo aktīvo pieprasījumu pieaugošā kv_len kopējā ietekme ir daudz mazāk izteikta. Par to ir jāmaksā ar laiku līdz pirmajam tokenam (TTFT) pieaugumu, jo KV kešatmiņas pārsūtīšana tīklā var aizņemt no desmitiem līdz simtiem milisekunžu.

KV Messenger

Uzņēmumā Perplexity mūsu implementācija sadalītajai sagatavošanai un dekodēšanai ir balstīta uz KV ziņojumapmaiņas rīku (messenger), kas mijiedarbojas ar LLM dzinēju, lai organizētu KV kešatmiņas pārsūtīšanu no sagatavošanas mezgliem uz dekodētāja mezgliem caur tīklu. Sagatavošanas pusē ziņojumapmaiņas rīks pieņem pieprasījumus no dekodētāja mezgliem, nododot tos pakešu plānotājam un sekojot līdzi priekšējās caurlaides izpildei, lai pēc iespējas ar mazāku latentumu nosūtītu KV kešatmiņas. Dekodētāja pusē pēc neizdzēšamo lapu (pages) piešķiršanas ziņojumapmaiņas rīks bloķē pieprasījuma ieplānošanu dekodēšanai, līdz tiek saņemts paziņojums par KV kešatmiņas un dekodētāja konteksta pārsūtīšanas pabeigšanu.

Izvērst

Sagatavošanas sadalīšanai ir nepieciešami augstas caurlaidspējas un zema latentuma savienojumi, tāpēc mūsu implementācija ir pielāgota RDMA, atbalstot gan EFA, gan ConnectX tīkla saskarnes kartes (NIC). KV Messenger ir veidots uz libfabric, izmantojot mūsu fabric-lib ietvarus (wrappers), lai nodrošinātu augstāka līmeņa zema latentuma abstrakcijas virs attālinātās tiešās atmiņas piekļuves (RDMA) primitīviem, ieviešot efektīvu lapu un metadatu pārsūtīšanu, kā arī zema latentuma signalizāciju. Fonā fabric-lib koordinē GPU un tam tieši pievienotos NIC, lai kopētu datus no sagatavošanas mezgla uz dekodētāja mezglu.

Saņemšanas brīdī sagatavošanas mezgls piešķir atbilstošu avota KV lapu kopu un ieplāno pieprasījumu sagatavošanai, izmantojot savu lokālo dzinēju. Lai samazinātu latentumu, pārsūtīšana negaida priekšējo caurlaidi: tā vietā KV lapu kopēšana tiek uzsākta, tiklīdz modelis pabeidz KV kešatmiņas ierakstu pievienošanu KV kešatmiņai atsevišķiem slāņiem. Tā kā sagatavošanas pieprasījumus var sadalīt daļās (chunked), pakešu plānotājs pirms izpildes paziņo KV ziņojumapmaiņas rīkam par pašlaik ieplānotajām daļām. Lai atbalstītu CUDA grafikus, vienlaikus saglabājot spēju izsekot slāņiem, ziņojumapmaiņas rīks uztur īpašu pavediens, kas aptaujā skaitītāju, kurš tiek palielināts pēc uzmanības izvades projekcijas. Skaitītājs tiek uzturēts tikai vadošajā mezglā saslāņotā (sharded) vidē: lai gan KV kešatmiņas ieraksti ir derīgi pēc pievienošanas un pirms uzmanības, izvades projekcija tiek reducēta visos rangos, netieši tos sinhronizējot. Kad tiek novērotas izmaiņas skaitītājā, ziņojumapmaiņas rīks saņem paziņojumu un izsauc fabric-lib, lai uzsāktu slāņa pārsūtīšanu.

Izvērst

Pēc pēdējās daļas pārsūtīšanas pabeigšanas tiek nokopēti arī visi papildu metadati: spekulatīvā dekodēšana vai MTP prasa logits un slēpto stāvokļu pārvietošanu uz dekodētāju. Šī kopēšana tiek veikta arī, izmantojot RDMA, uz un no iepriekš piešķirtiem buferiem.

Pēc visu pēdējās daļas gaidīto pārsūtīšanu pabeigšanas sagatavošanas mezgls atbrīvo KV lapas un pabeidz pieprasījumu. Dekodētāja mezglam netiek tieši paziņots: tā vietā tas izmanto tūlītējus skaitītājus, lai sekotu līdzi pabeigto operāciju skaitam. RDMA operāciju skaits sagatavošanas pusē ir proporcionāls pārsūtīto lapu skaitam. Pēc zināmā lapu un konteksta kopiju skaita pabeigšanas fabric-lib izsauc KV ziņojumapmaiņas rīku, lai norādītu, ka pieprasījums ir gatavs dekodēšanai. Ziņojumapmaiņas rīks atbrīvo visu kontekstu un nodod pieprasījumu LLM dzinējam.

Sadalītas KV kešatmiņas pārsūtīšana

Ja sagatavošana un dekodēšana balstās uz tenzoru paralēlismu (TP) un sadala vai replicē KV kešatmiņas identiski, viens pārsūtīšanas dzinējs koordinē vairākas ierīces, lai nosūtītu un saņemtu visu repliku lapas. Lai varētu izmantot vienu ziņojumapmaiņas rīku un pārsūtīšanas dzinēju, neskatoties uz to, ka modeļa izpildītājs ir replicēts vairākās ierīcēs un procesos, tiek izmantoti cuMem un cuMemImportFromShareableHandle, lai piešķirtu ierīces atmiņu, kas nodrošina KV kešatmiņas, un kartētu to galvenajā procesā. Pārsūtīšanas dzinējs pārbauda mezgla topoloģiju, lai atrastu NIC un CPU tuvākajā NUMA mezglā, ko izmantot katras KV kešatmiņas daļas pārsūtīšanai.

Ja avota un galamērķa daļas sakrīt identiski, pārsūtīšana ir vienkārša, jo pastāv savstarpēja vien-pret-vienu atbilstība starp avota un galamērķa ierīcēm un lapām. Šajā situācijā sadalīšana netieši palīdz pārsūtīšanas latentumam: izmantojot vairāk GPU, var izmantot vairāk saistīto NIC, tuvojoties pilnam joslas platuma izmantojumam. Tomēr, ja pastāv neatbilstība, pārsūtīšanas dzinējam ir jāsadala vai jārekonstruē lapas atkarībā no attiecības starp avota un galamērķa daļām.

Izvērst

Ja sagatavošanas dzinējs sadala KV kešatmiņu starp vairākām ierīcēm, pilnas lapas tiek rekonstruētas dekodētājā, nosūtot atbilstošās puses no sagatavošanas ierīcēm. Ja dekodētājam ir vairāk daļu, tas saņem lapas no vairākiem avotiem. Dekodētājam ir jāzin sagatavošanas shēma, lai varētu aprēķināt RDMA rakstīšanas operāciju skaitu, ko tas sagaida saņemt. Ja tiek izmantota replikācija, sagatavošanas dzinējs sagrupē ierīces repliku kopās, kas pašas par sevi replicē pilnu KV kešatmiņu. Galamērķa repliku kopām tiek nejauši piešķirta viena no avota kopām, lai izmantotu visas pieejamās ierīces RDMA rakstīšanas uzsākšanai.

Izvērst

Sadalīta pārsūtīšana prasa nelielas izmaiņas KV kešatmiņās. Pēc noklusējuma FlashInfer balstās uz NHD izkārtojumu, kas sakārto tokenus lapā galviņu (heads) ietvaros. Tā kā kešatmiņas visdrīzāk ir sadalītas pa uzmanības galviņu skaitu, tas rada pārtraukumus galviņas ietvaros. RDMA pārsūtīšana netieši neatbalsta soļotas rakstīšanas (strided writes), tāpēc pārsūtīšanas veikšanai ir nepieciešama viena operācija uz katru galviņu. Tā vietā, lai samazinātu mijiedarbības reižu skaitu ar libfabric, mēs organizējam KV kešatmiņas, izmantojot HND izkārtojumu, kas novieto galviņas dimensiju pirms tokenu skaita. Tas nodrošina nepārtrauktību, ļaujot nokopēt lapu ar vienu rakstīšanas operāciju.

Spekulatīvā dekodēšana

Spekulatīvā dekodēšana prasa nelielas izmaiņas sadalītajā sagatavošanā un dekodēšanā. Mūsu implementācijā sagatavošanas mezgliem nav atļauts veikt tokenu paraugu ņemšanu (sampling). Tā kā Perplexity Sonar modeļi atbalsta strukturētu izvadi, mēs nevēlamies uzņemties shēmas procesora implementāciju sinhronizācijas sarežģītību starp sagatavošanas un dekodēšanas mezgliem. MTP un spekulatīvās dekodēšanas mehānismos melnraksta modeļa sagatavošana līdz pēdējam tokenam ietver tokenu paraugu ņemšanu no mērķa modeļa.

Izvērst

Lai apietu šīs problēmas, sagatavošana neietver ievades secības pēdējo tokenu. Tā vietā slēptie stāvokļi vai logits no sagatavošanas pirms pēdējā tokena tiek pārsūtīti, un nākamajā solī dekodētājā tas tiek uzskatīts par dekodēšanas tokenu. Lai gan tas nedaudzpalielina latentumu, jo pēc sagatavošanas ir jāveic pilns dekodēšanas solis, lai izsniegtu pirmo tokenu, ieviešanas sarežģītība ir ievērojami samazināta.

Sadalītas izvietošanas

Mēs esam ieviesuši vai eksperimentējuši ar vairākām sadalītām konfigurācijām ar dažādiem modeļiem, lai atbalstītu vai nu ražošanas trafiku, vai iekšējās novērtēšanas darba slodzes. Pamatojoties uz modeļu izmēru un uzmanības mehānismu, mēs izvēlējāmies piemērotas sagatavošanas un dekodēšanas mezglu sadalīšanas shēmas, lai vislabāk izmantotu GPU.

DeepSeek-R1

Izmantojot DeepSeek, mēs apsvērām gan tenzoru paralēluma (TP), gan datu paralēluma (DP) izvietošanu. Kā apspriests iepriekšējos emuāra ierakstos, TP izvietošana nodrošina labāku latentumu uz zemākas caurlaidspējas rēķina, pieprasot vairāk GPU smagas trafikas apkalpošanai. DP izvietošana mērogojas daudz labāk ar slodzi, tomēr tās maksimālā caurlaidspēja ir zemāka ierīču vai mezglu starpsavienojumu izmaksu dēļ.

DeepSeek balstās uz vairāku galviņu latentās uzmanības (Multi-Head Latent Attention) mehānismu, saspiežot KV kešatmiņas. Tā kā visas KV galviņas tiek saspiestas vienā latentā vektorā, TP nevar sadalīt KV kešatmiņas, jo tam tā vietā ir jāreplicē latenti vektori visos rangos. Sadalīšana notiek pēc atspiešanas, jo katrs rangs var izgūt dažādas galviņas no tās pašas latentās reprezentācijas. Līdz ar to visas KV kešatmiņas daļas ir identiskas gan sagatavošanas, gan dekodēšanas mezglos.

Izmantojot mezgla iekšējo TP iestatījumu, gan sagatavošanas, gan dekodēšanas mezgli ir sadalīti identiski. Pārsūtīšana tiek izsūtīta no visiem rangiem, lai pilnībā izmantotu visus pieejamos NIC. Tomēr ar DP izvietošanu, kur TP ranga izmērs ir mazāks vai katram DP rangam ir piešķirts viens GPU, jebkura sagatavošanas ierīce, kurā ir KV kešatmiņas replicēta kopija, var to nosūtīt. Lai sabalansētu pieprasījumus starp visiem pieejamajiem NIC, mēs nejauši izvēlamies GPU un NIC, lai nosūtītu KV kešatmiņu no sagatavošanas mezgla uz dekodētāja mezglu.

Ar jauktu sagatavošanu un dekodēšanu mūsu R1 izvietošanai bija grūti stabili pārsniegt 50 TPS biežu sagatavošanas pārtraukumu dēļ simtu milisekuņu robežās. Turpretī, atdalot sagatavošanu, mēs ieguvām aptuveni 100 ms sodu TTFT katram pieprasījumam, taču viens sagatavošanas mezgls varēja uzturēt konsekventus pakešu izmērus uz 3 dekodētāja mezgliem, nodrošinot caurlaidspēju, kas pārsniedz 90 TPS, vienlaikus apstrādājot aptuveni 1 QPS slodzi uz katru dekodētāja mezglu. Izmantojot datu paralēlas izvietošanas, TPS bija nedaudz zemāks – ap 50, taču instances spēja apstrādāt 1 QPS slodzi uz katru rangu, ar 8 rangiem uz vienu mezglu.

Qwen3-Coder

Šis 480B modelis izmanto grupēto vaicājumu uzmanību (GQA), tāpēc uzmanību var viegli sadalīt un gūt labumu no tenzoru paralēluma, neupurējot atmiņu KV kešatmiņām. Līdz ar to mēs varējām sadalīt modeli starp 8 GPU gan sagatavošanai, gan dekodēšanai, savienojot apmēram 3 dekodētāja mezglus ar vienu sagatavošanas mezglu. Tā kā uzmanība ir sadalīta, mēs paļaujamies uz HND KV kešatmiņas izkārtojumu, lai sadalītu sagatavošanas un dekodēšanas KV kešatmiņas, savienojot sagatavošanas rangus ar dekodētāja rangiem un pilnībā izmantojot visus NIC, lai paralēli pārsūtītu daļas.

Citāts

@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}
}