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

Елегантен инструмент за справяне с нововъзникващите комуникационни модели при LLM

АвториPerplexity Research

LLM системите традиционно разчитат на колективни комуникации за структурирани паралелни модели при обучение и инференс. Въпреки това нововъзникващите работни натоварвания, като дезагрегиран инференс, маршрутизиране на Mixture-of-Experts (MoE) и асинхронно обучение за фина настройка с подкрепящо обучение, изискват гъвкавостта на точка-до-точка комуникацията.

Нашата най-нова изследователска статия представя 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: Нисколатентни операции за разпределение и комбиниране на MoE с помощта на разпръснати (scatter) и бариерни примитиви, постигащи най-съвременна латентност на ConnectX-7 и първата жизнеспособна реализация, съвместима с AWS EFA.

Тази статия се фокусира върху самия TransferEngine – комуникационната библиотека, която захранва тези системи. Прочетете пълната изследователска статия в arXiv и разгледайте кода в GitHub.

Мотивация

Колективните комуникации се отличават при структурирани паралелни модели като тензорен паралелизъм и паралелизъм на данните, но те налагат ограничения, които ограничават тяхната приложимост към нововъзникващи работни натоварвания:

Фиксирано членство: Колективите изискват установяване на глобална комуникационна среда в момента на инициализация, със синхронизирана настройка сред всички участници. За дезагрегирания инференс това е проблематично, тъй като броят на екземплярите за prefill и decode трябва да се мащабира динамично в отговор на трафик моделите.

Еднообразни тензорни форми: Колективите приемат, че всички участници обменят тензори с еднакви форми и типове данни. При MoE маршрутизирането броят на токените, изпратени до всеки експерт, зависи от данните и се установява по време на изпълнение, което принуждава плътните all-to-all колективи да се добавят с подложки до най-лошите размери на буфера.

Подредба на операциите: Колективните API налагат синхронизирани семантики, дори когато приложенията трябва само да знаят, че даден набор от операции е завършен. Например при прехвърлянето на KvCache всички страници от екземпляра за prefill към екземпляра за decode могат да пристигнат в произволен ред. По подобен начин при актуализациите на теглата в RL всички тензори с тегла могат да бъдат прехвърлени независимо без ограничения за подредбата.

Базираната на RDMA точка-до-точка комуникация естествено адресира тези ограничения, но три бариери възпрепятстват нейното приемане:

Ограничения на библиотеките: Въпреки че torch.distributed и NCCL предоставят API за точка-до-точка като Send/Recv, тези операции все пак наследяват колективни семантики и ограничения.

Хардуерно разнообразие: NVIDIA ConnectX NIC с надеждна връзка (Reliable Connection - RC транспорт) са най-разпространената RDMA настройка, гарантираща доставка в ред. Облачните доставчици като AWS използват различен хардуер – Elastic Fabric Adapter (EFA) с транспорт Scalable Reliable Datagram (SRD), който доставя съобщения извънредно (out-of-order). EFA също така изисква агрегиране на четири 100 Gbps NIC (или два 200 Gbps NIC), за да се постигне пропускателна способност от 400 Gbps.

Софтуерна преносимост: Съществуващите RDMA библиотеки страдат от обвързване с доставчик (vendor lock-in). DeepEP изисква инициирано от GPU RDMA (IBGDA), достъпно само на ConnectX. NVSHMEM проявява сериозно влошаване на производителността на EFA. Библиотеки като Mooncake и NIXL нямат поддръжка за EFA или все още са в ранен стадий на развитие. В резултат на това не съществува жизнеспособно крос-доставчик решение за LLM инференс.

TransferEngine

TransferEngine адресира тези бариери чрез четири основни възможности:

Общи транспортни предположения: Вместо да се насочва към конкретен NIC, TransferEngine приема само че основният транспорт е надежден, но ненареден, т.е. сечението на гаранциите, предоставени както от RC, така и от SRD транспортите.

ImmCounter синхронизация: Новаторски примитив за нотификация за завършване, който работи без да разчита на подредба на съобщенията, позволявайки на един и същ код да се изпълнява както на подредени (RC), така и на ненаредени (SRD) транспорти.

Поддръжка на много NIC (Multi-NIC): Прозрачно управление на множество NIC на GPU, автоматично раздробяване на трансферите за максимизиране на пропускателната способност.

Минимална API повърхност: Малък набор от операции, достатъчен за изграждане на сложни комуникационни модели.

Системен дизайн на TransferEngine, показващ пътищата на заявките и обратните извиквания, CPU работници, UVM буфери, GPU и NIC в два GPU/NUMA домейна.

TransferEngine стартира по една работна тема на GPU. Всяка тема управлява DomainGroup, която координира от 1 до 4 RDMA NIC в зависимост от хардуерната платформа (единен ConnectX-7 осигурява 400 Gbps, докато EFA изисква четири 100 Gbps адаптера или два 200 Gbps адаптера). В рамките на DomainGroup всеки Domain е хардуерно специфичен и управлява един NIC, обработвайки управлението на опашките от двойки (queue pairs), подаването на работа и проверката за завършване.

Всеки екземпляр на TransferEngine излага един главен адрес за идентификация и откриване на партньори. Отдалечените партньори обменят NetAddr структури, които кодират мрежови адреси и могат да бъдат серилизирани за извънлентова комуникация. Всички партньори в системата трябва да използват еднакъв брой NIC на GPU. Тази еднообразна конфигурация позволява на TransferEngine да раздробява трансферите интелигентно: при започване на трансфер както източникът, така и дестинацията знаят точно кои NIC са налични, което позволява автоматично балансиране на натоварването между всички адаптери.

API дизайн

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 операциите изискват регионите на паметта да бъдат регистрирани в NIC преди прехвърлянето на данни. Функцията reg_mr приема указател, дължина и устройство, връщайки два дескриптора:

MrHandle: Локален дескриптор, използван като източник за трансфери

MrDesc: Сериализируем дескриптор, който може да бъде обменян с отдалечени партньори, за да им позволи да пишат в регистрираната памет

Структурата MrDesc съдържа списък с отдалечени мрежови адреси и отдалечен ключ (по един на NIC), които дават отдалечен достъп до региона на паметта. Този дизайн естествено се справя с конфигурации с множество NIC.

Двустранни Send/Recv

За малки съобщения и съобщения за контролен равнин стил RPC, TransferEngine предоставя submit_send и submit_recvs:

submit_send: Изпраща съобщението до отдалечен партньор, идентифициран от NetAddr. Съобщението се копира вътрешно, така че извикващият може незабавно да го използва повторно или да освободи буфера.

submit_recvs: Публикува пул от буфери за получаване. Когато пристигне съобщение, се извиква обратно извикване с получените данни. Буферът се връща автоматично в пула, след като обратната връзка приключи.

Едностранен запис (One-Sided Write)

За масови трансфери на данни TransferEngine предоставя операции за едностранен запис без копиране (zero-copy):

submit_single_write: Прехвърля съседен регион от източника към дестинацията

submit_paged_writes: Прехвърля несъседни страници, зададени от индекси, стъпки (strides) и отмествания (offsets)

И двете операции приемат незадължителна 32-битова незабавна стойност, която се доставя на приемника атомно при завършване.

Системата автоматично раздробява големи трансфери през множество NIC, за да максимизира пропускателната способност.

Партньорски групи и масови операции

За операции, насочени към множество партньори (често срещани при MoE маршрутизиране), TransferEngine предоставя оптимизирани примитиви:

add_peer_group: Предварително регистрира набор от адреси на партньори за бързо търсене

submit_scatter: Изпраща различни филийки от данни до всеки партньор в групата

submit_barrier: Изпраща нотификация само с незабавна стойност до всички партньори (трансфер с нула байта)

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

UVM Watcher: Сигнализиране от GPU към CPU

Ключово предизвикателство в LLM системите е координирането на RDMA трансферите с напредъка на GPU. В идеалния случай трансферът на данни трябва да започне веднага щом GPU произведе изход, но анкетирането (polling) на GPU от CPU е неефективно.

TransferEngine адресира това с UVM Watchers: унифицирани места във виртуалната памет, които могат да бъдат актуализирани от GPU ядра (дори в рамките на CUDA графики) и ефективно анкетирани от CPU с помощта на GDRCopy. Когато се разпредели UVM watcher, се регистрира обратна връзка, която се извиква всеки път, когато GPU актуализира стойността.

Това позволява модели като прехвърляне на KvCache слой по слой: след завършването на всеки трансформерен слой GPU увеличава watcher-а, задействайки RDMA трансфер на KvCache на този слой, докато следващите слоеве продължават да изчисляват.

Нотификация за завършване

TransferEngine предоставя механизми за нотификация за завършване както за инициаторната страна (извикващият submit_write), така и за приемната страна (целта на Write операциите).

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

Нотификация от страна на приемника: Ключовата иновация е ImmCounter – компонент, който следи броячи, индексирани от 32-битови незабавни числа. Когато операция за запис носи незабавна стойност, ImmCounter на приемника се увеличава атомно при доставката на данните. Функцията expect_imm_count позволява на приемника да регистрира обратна връзка, която се задейства, когато специфична незабавна стойност достигне целеви брой. Броячите могат да бъдат прозрачно синхронизирани с GPU чрез GDRCopy, анкетирани директно или наблюдавани чрез обратни извиквания.

Този механизъм от страната на приемника осигурява надеждна синхронизация без изискване за подредба на съобщенията. Работи еднакво добре на ConnectX NIC (които гарантират доставка в ред чрез RC транспорт) и EFA NIC (които доставят съобщения извънредно чрез SRD транспорт).

Реализация

TransferEngine е реализиран на Rust, като внимателно оптимизира алокациите, нишките и синхронизацията, за да минимизира латентността и да постигне висока пропускателна способност. Системата стартира по една работна тема на всяка DomainGroup, като всяка е привързана към CPU ядро на NUMA възлите, към които са свързани устройствата, минимизирайки латентността както на планирането, така и на достъпа до паметта. Структурите от данни, специфични за даден домейн, се разпределят след привързването, за да се гарантира, че паметта е запазена в правилния NUMA възъл. Една работна тема управлява до 4 Domain (всяка управляваща един NIC), докато друга специализирана тема анкетира GPU за актуализиране на UVM watchers. Междутемната комуникация се извършва чрез опашки без заключване (lock-free queues).

API препраща заявките към DomainGroup, обслужваща съответното устройство, като първата също обслужва хоста. В стегнат цикъл работната тема на домейна проверява за нова работа, като дава приоритет на подаването на нови заявки. В зависимост от хардуера и конфигурацията, заявките се раздробяват и балансират между наличните Domain. Първият запис на композитна заявка се публикува незабавно в изходната опашка на NIC. След като новите заявки бъдат изчерпани, работникът премива към напредване на чакащите операции, публикувайки записи за запълване на хардуерния пайплайн. Впоследствие опашките за завършване се анкетират за заявки за завършени трансфери и увеличения на незабавния брояч. Събитията се агрегират за доставяне на нотификации за всеки трансфер, прехвърляйки трансфера към специализирана тема за обратна връзка, споделена от всички групи.

Раздробяването в рамките на DomainGroup е гъвкаво. Трансферите могат да се насочват към конкретни NIC по индекс. Един единствен запис може да бъде разделен. Пейджираните трансфери, scatter и бариерните операции, които се транслират в множество записи, могат да се раздробяват през всички NIC.

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

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

AWS EFA: Реализацията на EFA използва libfabric, управлявайки фабричен домейн на NIC в рамките на DomainGroup. Ключовите решения за реализация и оптимизации включват:

Записи с нулев размер: EFA се отклонява от RDMA спецификацията, изисквайки валидни целеви дескриптори дори за записи с нулев размер. Следователно TransferEngine налага валидни дескриптори за всички трансфери.

Шаблони за работни заявки: Предварително попълва и запазва общите полета на libfabric структури преди публикуване, намалявайки овърхеда на операция за масови трансфери и операции на партньорски групи.

NVIDIA ConnectX-7: Реализацията на ConnectX-7 използва libibverbs с RC транспорт. Ключовите решения за реализация и оптимизации включват:

UD опашки от двойки за ръкостискания: Всеки Domain използва опашка от двойки Unreliable Datagram (UD) за обмен на ръкостискания за RC връзка.

Отделни RC опашки от двойки: Създава 2 RC опашки от двойки на партньор – една за двустранни Send/Recv операции и друга за едностранни Write и WriteImm операции. Това разделяне е необходимо, защото както завършванията на Recv, така и на WriteImm консумират работни заявки в ред на публикуване, позволявайки семантика на Recv от високо ниво, като същевременно поддържат WriteImm без намеса.

Шаблони за работни заявки: Предварително попълва общите полета на структурите за работни заявки за масови трансфери и партньорски групи.

Веригообразуване на работни заявки: Свързва до 4 работни заявки чрез next указателя на ibv_send_wr, намалявайки броя на звъненията на звънеца (doorbell rings) към NIC.

Свободна подредба (Relaxed ordering): Активира IBV_ACCESS_RELAXED_ORDERING, за да позволи извънредни PCIe транзакции между NIC и паметта на 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}
}