RDMA комуникација тачка-у-тачку за LLM системе

LLM системи су се традиционално ослањале на колективне комуникације за структуриране паралелне обрасце у фази обуке и закључивања. Међутим, нове радне оптерећености, као што су дезагрегирано закључивање, рутирање смеше стручњака (MoE) и асинхроно подешавање кроз реинфорсмент учење (RL), захтевају флексибилност тачка-у-тачку комуникације.

Наш најновији истраживачки рад представља TransferEngine, преносиву комуникациону библиотеку засновану на RDMA технологији која ради на различитим мрежним адаптерима (NIC), укључујући NVIDIA ConnectX-7 и AWS Elastic Fabric Adapter (EFA). Демонстрирамо њену ефикасност у три продукциона система у компанији Perplexity AI:

Дезагрегирани Prefill и Decode: Пренос KvCache података по слојевима коришћењем пагованог уписа (paged writes) и UVM надгледача, што омогућава да се prefill и decode инстанце независно скалирају.

Пренос тежина за RL накнаду обуку за мање од 2 секунде: Невероватно брза синхронизација тежина између чворова за обуку и закључивање.

Омогућавање модела са билион параметара на AWS EFA: Операције диспечерског слања и комбиновања (dispatch и combine) са малом латенцијом за MoE, коришћењем scatter и barrier примитива, чиме се постиже латенција највишег нивоа на ConnectX-7 и прва одржива имплементација компатибилна са AWS EFA.

Овај чланак се фокусира на сам TransferEngine, комуникациону библиотеку која покреће ове системе. Прочитајте комплетан истраживачки рад на платформи arXiv и истражите код на GitHub платформи.

Колективне комуникације су одличне за структуриране паралелне обрасце као што су паралелизам тензора и паралелизам података, али намећу ограничења која ограничавају њихову применљивост на нове радне оптерећености:

Фиксно чланство: Колективне комуникације захтевају успостављање глобалног комуникационог света у тренутку иницијализације, уз синхронизовано подешавање свих учесника. Ово представља проблем код дезагрегираног закључивања пошто број prefill и decode инстанци мора динамички да се скалира као одговор на обрасце саобраћаја.

Униформни облици тензора: Колективне комуникације претпостављају да сви учесници размјењују тензоре идентичних облика и типова података. У MoE рутирању, број токена послатих сваком стручњаку зависи од података и открива се у вријеме извршавања, што приморава густе all-to-all колективне комуникације да се допуне (pad) до величина бафера за најгори сценарио.

Редослед операција: Колективни API-ји намећу синхронизовану семантику, чак и када апликације треба само да знају да је скуп операција завршен. На пример, у преносу KvCache-а, све странице од prefill инстанце до decode инстанце могу стићи било којим редоследом. Слично томе, код ажурирања тежина у RL-у, сви тензори тежина могу се преносити независно без ограничења редоследа.

RDMA комуникација тачка-у-тачку природно рјешава ова ограничења, али су три препреке спречиле њено усвајање:

Ограничења библиотеке: Иако torch.distributed и NCCL пружају API-је тачка-у-тачку попут Send/Recv, ове операције и даље наслеђују семантику и ограничења колективних комуникација.

Хардверска разноврсност: NVIDIA ConnectX мрежни адаптери са транспортним механизмом Reliable Connection (RC) представљају најчешћу RDMA поставку, која гарантује испоруку у одређеном редоследу. Пружаоци облачних услуга попут AWS користе другачији хардвер, Elastic Fabric Adapter (EFA) са Scalable Reliable Datagram (SRD) транспортним механизмом који поруке испоручује ван редоследа. EFA такође захтева агрегацију четири мрежна адаптера од 100 Gbps (или два од 200 Gbps) да би се постигла пропусни опсег од 400 Gbps.

Софтверска преносивост: Постојеће RDMA библиотеке пате од везаности за једног добављача (vendor lock-in). DeepEP захтева RDMA инициран од стране GPU-а (IBGDA) који је доступан само на ConnectX-у. NVSHMEM показује озбиљно погоршање перформанси на EFA-у. Библиотеке попут Mooncake и NIXL немају подршку за EFA или су у раној фази развоја. Због тога не постоји одрживо рјешење за више добављача за LLM закључивање.

TransferEngine рјешава ове препреке кроз четири основне могућности:

Униформне претпоставке о транспорту: Умјесто да циља одређени мрежни адаптер, TransferEngine претпоставља само да је основни транспорт поуздан али неуређен, тј. да представља пресек гаранција које пружају и RC и SRD транспорти.

ImmCounter синхронизација: Нови примитив за обавештење о завршетку који ради без ослањања на редослед порука, омогућавајући да се исти код извршава и на уређеним (RC) и на неуређеним (SRD) транспортима.

Подршка за више мрежних адаптера: Транспарентно управљање више мрежних адаптера по једном GPU-у, аутоматским дијељењем преноса на делове (sharding) ради максимизирања пропусног опсега.

Минимална површина API-ја: Мали скуп операција довољних за изградњу сложених комуникационих образаца.

Dizajn sistema TransferEngine, koji prikazuje putanje zahtjeva i povratnih poziva, CPU radnike, UVM baferе, GPU-ove i NIC-ove preko dva GPU/NUMA domena.

TransferEngine покреће један радни тред по GPU-у. Сваки радни тред управља са DomainGroup која координира од 1 до 4 RDMA мрежна адаптера, у зависности од хардверске платформе (један ConnectX-7 обезбеђује 400 Gbps, док EFA захтева четири адаптера од 100 Gbps или два адаптера од 200 Gbps). Унутар DomainGroup, сваки Domain је специфічан за хардвер и управља једним мрежним адаптером, руководећи управљањем паровима редова (queue pair), подношењем посла и анкетирањем завршетка.

Свака инстанца TransferEngine-а излаже једну главну адресу за идентификацију и откривање парњака. Удаљени парњаци размјењују NetAddr структуре које кодирају мрежне адресе и могу се серијализовати за ванпојасну комуникацију. Сви парњаци у систему морају користити исти број мрежних адаптера по GPU-у. Ова униформна конфигурација омогућава TransferEngine-u да паметно дијели преносе: при покретању преноса, и извор и одредиште тачно знају који су мрежни адаптери доступни, што омогућава аутоматско балансирање оптерећења на свим адаптерима.

TransferEngine излаже минимални API који апстрахује сложеност основних RDMA интерфејса.

rust

#[serde] struct NetAddr(Bytes);
#[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> }
struct MrHandle(NonNull<c_void>); type Offset = u64;
struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset }
struct PeerGroupHandle(u64);
struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)}
enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) }

trait TransferEngine {
  fn main_address() -> NetAddr;
  // Memory Region Management
  fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc);
  // Two-sided Send/Recv
  fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ());
  fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ());
  // One-sided Write
  fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ());
  fn submit_single_write(len: u64, imm: Option<u32>
      src: (MrHandle, Offset), dst: (MrDesc, Offset), OnDone);
  fn submit_paged_writes(page_len: u64, imm: Option<u32>,
      src: (MrHandle, Pages), dst: (MrDesc, Pages), OnDone);
  // One-sided Write to a group of peers
  fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle;
  fn submit_scatter(h: Option<PeerGroupHandle>, OnDone,
      imm: Option<u32>, src: MrHandle, dst: Vec<ScatterDst>);
  fn submit_barrier(h: Option<PeerGroupHandle>, OnDone,
      imm: u32, dst: Vec<MrDesc>);
  // Watcher for CPU-GPU synchronization
  fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>;
}

API се састоји од неколико кључних компоненти:

RDMA операције захтевају да меморијски региони буду регистровани са мрежним адаптером прије преноса података. Функција reg_mr узима показивач, дужину и уређај, враћајући двије ручке:

MrHandle: Локална ручка која се користи као извор за преносе

MrDesc: Серијализовани дескриптор који се може размијенити са удаљеним парњацима како би им се омогућило да уписују у регистровану меморију

Структура MrDesc садржи листу удаљених мрежних адреса и удаљених кључева (по један за сваки мрежни адаптер) који одобравају удаљени приступ меморијском региону. Овај дизајн природно управља поставкама са више мрежних адаптера.

За мале поруке и поруке контролне равни у стилу RPC-а, TransferEngine обезбеђује submit_send и submit_recvs:

submit_send: Шаље поруку удаљеном парњаку идентификованом са NetAddr. Порука се копира интерно, тако да позивалац може одмах поново користити или ослободити бафер.

submit_recvs: Поставља базен пријемних бафера. Када порука стигне, позива се повратни позив (callback) са примљеним подацима. Бафер се аутоматски враћа у базен након што се повратни позив заврши.

За масовне преносе података, TransferEngine обезбеђује једносмерне операције уписа без копирања (zero-copy):

submit_single_write: Преноси континуирани регион од извора до одредишта

submit_paged_writes: Преноси неконтинуиране странице дефинисане индексима, корацима и помјерајима

Обје операције прихватају опционалну 32-битну тренутну вриједност која се испоручује примаоцу асинхроно по завршетку.

Машина аутоматски дијели велике преносе на више мрежних адаптера да би максимизовала пропусни опсег.

За операције циљања више парњака (често код MoE рутирања), TransferEngine обезбеђује оптимизоване примитиве:

add_peer_group: Претходно региструје скуп адреса парњака за брзо претраживање

submit_scatter: Шаље различите дијелове података сваком парњаку у групи

submit_barrier: Шаље искључиво тренутно обавештење свим парњацима (пренос од нула бајтова)

Ове операције су оптимизоване за малу латенцију кеширањем интерних меморијских структура и темплетирањем RDMA радних захтјева, чиме се избегавају репетитивни трошкови подешавања.

Кључни изазов у LLM системима је координација RDMA преноса са напретком GPU-а. Идеално би било да пренос података почне чим GPU произведе излаз, али је анкетирање GPU-а од стране CPU-а неефикасно.

TransferEngine рјешава ово помоћу UVM надгледача (UVM Watchers): унификованих локација виртуелне меморије које могу бити ажуриране од стране GPU језгара (чак и у оквиру CUDA графова) и ефикасно анкетиране од стране CPU-а помоћу GDRCopy-ја. Када је UVM надгледач алоциран, региструје се повратни позив који се покреће кад год GPU ажурира вриједност.

Ово омогућава обрасце попут преноса KvCache-а по слојевима: након што се заврши сваки трансформаторски слој, GPU увећава вриједност надгледача, покрећући RDMA пренос KvCache-а тог слоја док наредни слојеви настављају са рачунањем.

TransferEngine обезбеђује механизме обавештења о завршетку и за иницијаторску страну (позивалац функције submit_write) и за пријемну страну (циљ операција уписа).

Обавештење на страни иницијатора: Када се операција уписа поднесе, позивалац може одабрати да буде обавештен о завршетку путем функције повратног позива или атомске засташице (flag). Ово обавештење потврђује да су подаци успешно пренети из изворног бафера.

Обавештење на страни примаоца: Кључна иновација је ImmCounter, компонента која прати бројаче индексиране помоћу 32-битних тренутних бројева. Када операција уписа носи тренутну вриједност, ImmCounter примаоца се увећава атомски при испоруци података. Функција expect_imm_count омогућава примаоцу да региструје повратни позив који се активира када одређена тренутна вриједност достигне циљни број. Бројачи се могу транспарентно синхронизовати са GPU-ом преко GDRCopy-ја, директно анкетирати или надгледати кроз повратне позиве.

Овај механизам на страни примаоца пружа поуздану синхронизацију без захтијевања редоследа порука. Ради подједнако добро на ConnectX мрежним адаптерима (који гарантују испоруку у одређеном редоследу преко RC транспорта) и EFA мрежним адаптерима (који испоручују поруке ван редоследа преко SRD транспорта).

TransferEngine је имплементиран у Rust језику, уз пажљиву оптимизацију алокација, тредова и синхронизације како би се минимизирала латенција и постигао висок проток. Машина покреће један радни тред по DomainGroup, при чему је сваки везан за CPU језгро на NUMA чвору за који су уређаји придружени, минимизирајући латенцију распоређивања и приступа меморији. Меморијске структуре специфичне за домен алоцирају се након везивања како би се осигурало да је меморија резервисана у исправном NUMA чвору. Један радни тред управља са до 4 Domain-а (сваки управља једним мрежним адаптером), док други наменски тред анкетира GPU ради ажурирања UVM надгледача. Међутредна комуникација се обавља преко редова без закључавања (lock-free queues).

API прослеђује захтеве ка DomainGroup која опслужује одговарајући уређај, при чему први такође опслужује хост. У уској петљи, доменски радник анкетира нови посао, дајући приоритет подношењу нових захтева. У зависности од хардвера и конфигурације, захтеви се дијеле и балансирају по оптерећењу широм доступних Domain-а. Први упис сложеног захтева се одмах поставља у ред за слање мрежног адаптера. Када се нови захтеви исцрпе, радник наставља са напретком чекајућих операција, постављајући уписе како би попунио хардверски цевовод. Након тога, редови чекања завршетка се анкетирају да би се провјерили завршени преноси и тренутна повећања бројача. Догађаји се агрегирају да би се испоручила обавештења по преносу, препуштајући пренос наменском треду за повратне позиве који деле све групе.

Дијељење (sharding) унутар DomainGroup је флексибилно. Преноси могу циљати одређене мрежне адаптере по индексу. Један упис се може подијелити. Паговани преноси, scatter и barrier операције, које се све преводе у више уписа, могу се дијелити на све мрежне адаптере.

Domain-и унутар TransferEngine-а су специјализовани и оптимизовани за хардвер под њиховом контролом:

AWS EFA: EFA имплементација користи libfabric, управљајући доменом тканине (fabric domain) по мрежном адаптеру у оквиру DomainGroup. Кључне одлуке имплементације и оптимизације укључују:

Писања нулте величине: EFA одступа од RDMA спецификације тако што захтијева важеће одредишне дескрипторе чак и за писања нулте величине. Стога, TransferEngine примјењује важеће дескрипторе за све преносе.

Темплатирање радних захтјева: Унапријед попуњава и задржава заједничка поља libfabric структура прије слања, чиме се смањује оптерећење по операцији за масовне преносе и операције група пандана

NVIDIA ConnectX-7: ConnectX-7 имплементација користи libibverbs са RC транспортом. Кључне одлуке имплементације и оптимизације укључују:

UD парови редова за руковања: Сваки Domain користи пар редова Unreliable Datagram (UD) за размјену руковања (handshakes) RC везе.

Одвојени RC парови редова: Креира 2 RC пара редова по парњаку, један за двосмерне Send/Recv операције и други за једносмерне Write и WriteImm операције. Ово раздвајање је неопходно зато што и Recv и WriteImm завршетци троше радне захтјеве у редоследу постављања, омогућавајући Recv семантику високог нивоа уз подршку за WriteImm без мешања.

Темплетирање радних захтјева: Претходно попуњава заједничка поља структура радних захтјева за масовне преносе и групе парњака.

Увезивање радних захтјева: Повезује до 4 радна захтјева кроз next показивач структуре ibv_send_wr, смањујући број звона врата (doorbell rings) ка мрежном адаптеру.

Опуштени редослед: Омогућава IBV_ACCESS_RELAXED_ORDERING ради дозвољавања PCIe трансакција ван редоследа између мрежног адаптера и GPU меморије.

Циљањем пресека гаранција широм различитих RDMA транспорта и увођењем синхронизације независне од редоследа путем ImmCounter-а, TransferEngine пружа преносиву основу за комуникацију тачка-у-тачку у LLM системима. Библиотека је отвореног кода и доступна на GitHub-у, а подстичемо заједницу да је истражи и гради на основу ње за нове комуникационе обрасце у великим системима машинског учења.

BibTeX

@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{RDMA point-to-point communication for LLM systems}},
  date = {2025-11-05},
  url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}

Мотивација

TransferEngine

Дизајн API-ја

Регистрација меморије

Двосмерно слање/пријем (Send/Recv)

Једносмерни упис (One-Sided Write)

Групе парњака и масовне операције

UVM надгледач: Сигнализација са GPU-а ка CPU-у

Обавештење о завршетку

Имплементација

Хардверски специфичне оптимизације

Закључак

Цитат