Разделени презареждане и декодиране

За да генерира изходни токени от входен промпт, LLM инференсът е разделен на два етапа: презареждане (prefill) и декодиране. Презареждането работи върху входните токени, попълвайки KV кешове, преди да влезе в етапа на декодиране, който генерира токени един по един. Докато единична стъпка на декодиране обикновено отнема десетки милисекунди, презареждането отнема значително повече време. Ако се изпълнява на същите устройства, смесването на презареждането с декодирането влошава производителността на декодирането. В тази статия изследваме едно утвърдено решение под формата на разделени презареждане и декодиране, като ги изпълняваме на отделни устройства, за да максимизираме както пропускателната способност на презареждането, така и латентността на декодирането.

Производителност: Презареждане срещу декодиране

В типичен двигател за обслужване на LLM, планировчикът на пакети избира заявки за обработка във всяка стъпка на изпълнение на модела. Когато работи на едно устройство или въззел, както заявките за презареждане, така и тези за декодиране се пакетират заедно. Цената на вниманието (attention), която се агрегира по дължината на последователността, нараства както за презареждането, така и за декодирането, пропорционално на дължината на записа в KV кеша (kv_len). Заявките за декодиране обикновено препращат един единствен токен (qo_len=1) с минимална цена през други слоеве, които работят независимо върху токените на дадена последователност. Заявките за презареждане препращат хиляди или десетки хиляди токени със значителна цена през плътни слоеве (голям qo_len).

Латентността на едно forward преминаване се влияе по-силно от броя на независимите токени, преминаващи през плътни слоеве (qo_len), отколкото от броя на токените, извлечени от KV кеша по време на вниманието (kv_len). Вниманието може да се паралелизира едновременно по броя заявки и по kv_len, пропорционален на дължините на последователностите, постигайки добра степен на използване. Prefill е ограничен от изчисленията: при висок qo_len GEMM kernel-ите могат да заделят достатъчно блокове по измерението M, за да използват напълно изчислителните възможности на съвременните GPU. Decode е ограничен от паметта: поради обикновено ниските размери на батчовете броят на входовете по M обикновено е малък, достатъчен само за един блок. Макар Split-K GEMM kernel-ите да могат да подобрят използването на SM при ниски размери на token батчовете, кешовете и единиците за умножение на матрици обикновено остават слабо използвани.

Уголеми

Когато се смесят заедно, пакетите, съдържащи заявки за презареждане, водят до по-висока латентност при правото подаване, което се отразява негативно на пропускателната способност на декодирането за цялата инстанция. Макар че смесването на заявки за презареждане със заявки за декодиране или използването на сегментирано презареждане може леко да подобри производителността на декодирането, е трудно да се поддържа достатъчна пропускателна способност за презареждане, за да се обработват достатъчно заявки на дадена инстанция и да се максимизира пропускателната способност на декодирането. В случай на големи модели с типични изходни дължини, за да се поддържа голям размер на пакета за декодиране, презареждането трябва да се извършва достатъчно често, така че значително да влоши средната латентност и да причини накъсване на изхода.

Уголеми

Уголеми

Тези проблеми могат да бъдат решени чрез използване на отделен набор от възли за извършване на презареждане и декодиране. Чрез свързване на въззел за презареждане с множество възли за декодиране могат да бъдат планирани достатъчно заявки за презареждане, за да се максимизира пропускателната способност и да се поддържа достатъчно голям брой едновременни заявки на възлите за декодиране, така че да се максимизира и пропускателната способност на декодирането. Възлите за презареждане попълват KV кешовете, които след това се прехвърлят към възлите за декодиране. Тъй като декодерите вече не е нужно да прекъсват работата си за презареждане, латентността става много по-детерминирана, тъй като цялостното въздействие от нарастващия kv_len на активните заявки е много по-малко изразено. Цената се заплаща под формата на увеличение на времето до първия токен (Time to First Token - TTFT), тъй като прехвърлянето на KV кешовете през мрежата може да отнеме от десетки до стотици милисекунди.

KV Messenger

В Perplexity нашата реализация на разделени презареждане и декодиране е изградена около KV messenger, който взаимодейства с LLM двигателя, за да оркестрира прехвърлянето на KV кешове от възлите за презареждане към възлите за декодиране през мрежа. От страна на презареждането, messenger приема заявки от възлите за декодиране, предавайки ги на планировчника на пакети и проследявайки изпълнението на правото подаване, за да изпрати KV кешовете с възможно най-малка латентност. От страна на декодирането, след като бъдат заделени страници, които не могат да бъдат изведени, messenger блокира планирането на заявката за декодиране, докато не бъде уведомен за завършването на прехвърлянето на KV кеша и контекста на декодера.

Уголеми

Разделянето на презареждането изисква връзки с висока пропускателна способност и ниска латентност, поради което нашата реализация е адаптирана за RDMA, поддържайки както EFA, така и мрежови интерфейсни карти (NIC) ConnectX. KV Messenger е изграден върху libfabric, използвайки нашите обвивки fabric-lib, за да осигури абстракции от по-високо ниво с ниска латентност върху примитивите за отдалечен директен достъп до паметта (RDMA), реализирайки ефективни прехвърляния на страници и метаданни, заедно със сигнализиране с ниска латентност. във фонов режим fabric-lib координира графичен процесор и директно свързаните с него NIC за копиране на данни от въззела за презареждане към въззела за декодиране.

При получаване, въззелът за презареждане заделя съответен набор от изходни KV страници и планира заявката за презареждане, използвайки своя локален двигател. За да се минимизира латентността, прехвърлянията не чакат правото подаване: вместо това копирането на KV страници се инициира веднага щом моделът завърши добавянето на записи в KV кеша за отделните слоеве. Тъй като заявките за презареждане могат да бъдат сегментирани, планировчикът на пакети уведомява KV messenger за текущо планираните сегменти преди изпълнението. За да поддържа графики на CUDA, като същевременно може да проследява слоеве, messenger поддържа специална тема, която опросва брояч, инкрементиран след изходната проекция на вниманието. Броячът се поддържа само на водещия въззел в шардирана среда: въпреки че записите в KV кеша са валидни след добавяне и преди вниманието, изходната проекция се редуцира между ранговете, като ги синхронизира имплицитно. След като бъде отчетена промяна в брояча, messenger получава известие и извиква fabric-lib, за да инициира прехвърлянето на слой.

Уголеми

След като прехвърлянето на последния сегмент завърши, всички допълнителни метаданни също се копират: спекулативното декодиране или MTP изискват логитите и скритите състояния да бъдат преместени към декодера. Тези копия също се изпълняват чрез RDMA, към и от предварително заделени буфери.

След завършване на всички чакащи прехвърляния на последния сегмент, въззелът за презареждане освобождава KV страниците и завършва заявката. Въззелът за декодиране не се уведомява изрично: вместо това той използва незабавни броячи, за да проследява броя на завършените операции. Броят на RDMA операциите от страната на презареждането е пропорционален на броя на прехвърлените страници. След завършване на известния брой копия на страници и контекст, fabric-lib извиква KV messenger, за да покаже, че заявката е готова за декодиране. Messenger освобождава целия контекст и предава заявката на LLM двигателя.

Шардирани прехвърляния на KV кеш

Ако презареждането и декодирането разчитат на тензорен паралелизъм (Tensor Parallelism - TP) и шардират или репликират KV кешовете идентично, един единствен двигател за прехвърляне координира множество устройства за изпращане и получаване на страниците на всички реплики. За да може да се използва един единствен messenger и двигател за прехвърляне, въпреки факта, че изпълнителят на модела е репликиран на множество устройства и процеси, се използват cuMem и cuMemImportFromShareableHandle за заделяне на паметта на устройството, поддържаща KV кешовете, и за мапирането ѝ в главния процес. Двигателят за прехвърляне инспектира топологията на въззела, за да намери NIC и CPU в най-близкия NUMA въззел, които да използва за прехвърлянето на всеки от резените на KV кеша.

Ако източникът и дестинацията са шардирани идентично, прехвърлянията са елементарни, тъй като има съответствие едно към едно между устройствата и страниците на източника и дестинацията. В тази ситуация шардирането имплицитно помага на латентността на прехвърлянето: чрез използване на повече графични процесори могат да бъдат използвани повече свързани NIC, достигайки по-близо до пълното оползотворяване на пропускателната способност. Въпреки това, ако има несъответствие, двигателят за прехвърляне трябва да раздели или реконструира страниците в зависимост от съотношението между изходния и целевия резен.

Уголеми

Ако prefiller-ът разделя KV кеша между повече устройства, пълните страници се реконструират на decoder-а чрез изпращане на съответните половини от устройствата на prefiller-а. Ако decoder-ът има повече shard-ове, той получава страници от множество източници. Decoder-ът трябва да знае схемата на sharding на prefiller-а, за да може да изчисли броя на RDMA записите, които се очаква да получи. Ако е включена репликация, prefiller-ът групира устройствата в replica sets, които репликират целия KV кеш вътре в себе си. Replica sets на дестинацията се разпределят на случаен принцип към едно от source sets, за да се използват всички налични устройства за иницииране на RDMA записи.

Уголеми

Шардираните прехвърляния изискват лека корекция на KV кешовете. По подразбиране FlashInfer разчита на оформлението NHD, което подрежда токените в рамките на страница в главите (heads). Тъй като кешовете най-вероятно са шардирани по броя на главите на вниманието, това създава прекъснатост в рамките на главата. RDMA прехвърлянията не поддържат имплицитно записи със стъпка (strided writes), което изисква по една операция на глава за извършване на прехвърлянето. Вместо това, за да намалим броя на взаимодействията с libfabric, организираме KV кешовете, използвайки оформлението HND, което поставя измерението на главата преди броя на токените. Това гарантира непрекъснатост, позволявайки една страница да бъде копирана с едно единствено записване.

Спекулативно декодиране

Спекулативното декодиране изисква леки корекции на разделеното презареждане и декодиране. В нашата реализация на възлите за презареждане не е позволено да семплират токени. Тъй като моделите Sonar на Perplexity поддържат структуриран изход, не искаме да поемаме сложността на синхронизирането на реализациите на процесора на схеми между модулите за презареждане и декодиране. В механизмите за MTP и спекулативно декодиране, предварителното зареждане на базовия модел до последния токен включва семплиране на токени от целевия модел.

Уголеми

За да се заобиколят тези проблеми, презареждането не включва последния токен от входната последователност. Вместо това скритите състояния или логити от презареждането преди последния токен се прехвърлят и то се третира като токен за декодиране в следващата стъпка на декодера. Въпреки че това леко увеличава латентността, тъй като след презареждането трябва да се извърши пълна стъпка на декодиране за издаване на първия токен, сложността на реализацията е значително намалена.

Разделени разпогръщания (Disaggregated Deployments)

Ние сме разположили или експериментирали с множество разделени конфигурации с различни модели, за да поддържаме или производствен трафик, или вътрешни работни натоварвания за оценка. Въз основа на размера и механизма на внимание на моделите, ние избрахме подходящи схеми за шардиране за възлите за презареждане и декодиране, за да оползотворим най-добре графичните процесори.

DeepSeek-R1

При DeepSeek обмислихме разполагане както с тензорен паралелизъм (TP), така и с паралелизъм на данните (Data-Parallel - DP). Както беше обсъдено в предишни публикации в блогове, TP разполаганията осигуряват по-добра латентност за сметка на по-ниска пропускателна способност, изискващи повече графични процесори за обслужване на тежък трафик. DP разполаганията мащабират много по-добре при натоварване, но тяхната пикова пропускателна способност е по-ниска поради цената на комуникацията между устройствата или между възлите.

DeepSeek разчита на латентно внимание с много глави (Multi-Head Latent Attention), компресиращо KV кешове. Тъй като всички KV глави са компресирани в един латентен вектор, TP не може да шардира KV кешовете, тъй като вместо това трябва да репликира латентните вектори на всички рангове. Шардирането се случва след декомпресията, тъй като всеки ранг може да извлече различни глави от същото латентно представяне. В резултат на това всички шардове на KV кеша са идентични както между шардовете за презареждане, така и между тези за декодиране.

При вътрешно-възлова TP настройка както модулите за презареждане, така и декодерите са шардирани идентично. Прехвърлянията се изпращат от всички рангове, за да се използват напълно всички налични NIC. Въпреки това при DP разполагане, където размерът на TP ранга е по-малък или на всеки DP ранг е назначен един единствен графичен процесор, всяко устройство за презареждане, което притежава репликирано копие на KV кеша, може да го изпрати. За да балансираме заявките между всички налични NIC, на случаен принцип избираме графичен процесор и NIC за изпращане на KV кеша от презареждането към декодера.

При смесено презареждане и декодиране нашето R1 разполагане се затрудняваше да надхвърли последователно 50 TPS поради чести прекъсвания на презареждането от порядъка на стотици милисекунди. За разлика от това, чрез разделяне на презареждането понесехме наказание от около 100 ms за TTFT за всяка заявка, но един единствен въззел за презареждане можеше да поддържа постоянни размери на пакетите на 3 въззела за декодиране, предоставяйки пропускателна способност над 90 TPS, докато обработва натоварване от около 1 QPS на въззел за декодиране. При разполагания с паралелизъм на данните TPS беше малко по-нисък – около 50, но инстанциите можеха да се справят с натоварване от 1 QPS на ранг, с 8 ранга на един въззел.

Qwen3-Coder

Този модел с 480 милиарда параметъра използва внимание с групови заявки (Grouped-Query Attention - GQA), така че вниманието може лесно да бъде шардирано и да се възползва от тензорния паралелизъм, без да се жертва памет за KV кешове. В резултат на това бихме могли да шардираме модела между 8 графични процесора както за презареждане, така и за декодиране, сдвоявайки около 3 въззела за декодиране с един единствен въззел за презареждане. Тъй като вниманието е шардирано, ние разчитаме на оформлението на KV кеша HND, за да шардираме KV кешовете за презареждане и декодиране, сдвоявайки ранговете за презареждане с ранговете за декодиране и оползотворявайки напълно всички NIC за паралелно прехвърляне на резени.

Цитат

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