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

Системы LLM традиционно опирались на коллективные коммуникации для структурированных параллельных шаблонов при обучении и выводе. Однако новые рабочие нагрузки, такие как дезагрегированный вывод, маршрутизация Mixture-of-Experts (MoE) и асинхронная донастройка с обучением с подкреплением (RL), требуют гибкости коммуникации «точка-точка».

В нашей последней исследовательской работе представлена TransferEngine, переносимая коммуникационная библиотека на базе RDMA, которая работает с различными сетевыми интерфейсными контроллерами (NIC), включая NVIDIA ConnectX-7 и AWS Elastic Fabric Adapter (EFA). Мы демонстрируем её эффективность в трёх производственных системах в Perplexity AI:

Дезагрегированные префиллинг и декодирование: послойная передача KvCache с использованием постраничной записи и UVM-наблюдателей, что позволяет экземплярам префиллинга и декодирования масштабироваться независимо.

Перенос весов для пост-обучения с подкреплением (RL) менее чем за 2 секунды: сверхбыстрая синхронизация весов между узлами обучения и вывода.

Поддержка моделей с триллионами параметров на AWS EFA: операции диспетчеризации и объединения MoE с низкой задержкой с использованием примитивов scatter и barrier, что позволяет достичь передовой задержки на ConnectX-7 и первой жизнеспособной реализации, совместимой с AWS EFA.

В этой статье рассматривается сама TransferEngine — коммуникационная библиотека, обеспечивающая работу этих систем. Ознакомьтесь с полной исследовательской работой на arXiv и изучите код на GitHub.

Коллективные коммуникации эффективны для структурированных параллельных шаблонов, таких как параллелизм тензоров и параллелизм данных, но они накладывают ограничения, ограничивающие их применимость к новым рабочим нагрузкам:

Фиксированный состав участников: коллективные операции требуют создания глобальной коммуникационной среды во время инициализации с синхронизированной настройкой всех участников. Для дезагрегированного вывода это проблематично, так как количество экземпляров префиллинга и декодирования должно динамически масштабироваться в ответ на структуру трафика.

Унифицированные формы тензоров: коллективные операции предполагают, что все участники обмениваются тензорами с одинаковыми формами и типами данных. В маршрутизации MoE количество токенов, отправляемых каждому эксперту, зависит от данных и определяется во время выполнения, что вынуждает плотные операции all-to-all дополнять буферы до размеров наихудшего случая.

Порядок операций: API коллективных операций обеспечивают синхронизированную семантику, даже если приложениям нужно только знать, что набор операций завершён. Например, при передаче KvCache все страницы от экземпляра префиллинга к экземпляру декодирования могут поступать в любом порядке. Аналогично, при обновлении весов RL все тензоры весов могут передаваться независимо без ограничений по порядку.

Коммуникация «точка-точка» на базе RDMA естественным образом решает эти проблемы, однако её внедрению препятствуют три барьера:

Ограничения библиотек: хотя torch.distributed и NCCL предоставляют API «точка-точка», такие как Send/Recv, эти операции всё равно наследуют коллективную семантику и ограничения.

Аппаратное разнообразие: NIC NVIDIA ConnectX с транспортом Reliable Connection (RC) являются наиболее распространённой настройкой RDMA, гарантирующей доставку в порядке очереди. Облачные провайдеры, такие как AWS, используют другое оборудование: Elastic Fabric Adapter (EFA) с транспортом Scalable Reliable Datagram (SRD) доставляет сообщения не по порядку. EFA также требует агрегации четырёх NIC по 100 Гбит/с (или двух по 200 Гбит/с) для достижения пропускной способности 400 Гбит/с.

Переносимость программного обеспечения: существующие библиотеки RDMA страдают от привязки к поставщику. DeepEP требует GPU-инициируемого RDMA (IBGDA), доступного только на ConnectX. NVSHMEM демонстрирует серьёзное ухудшение производительности на EFA. Библиотеки, такие как Mooncake и NIXL, не имеют поддержки EFA или находятся на ранней стадии разработки. В результате не существует жизнеспособного кросс-провайдерного решения для вывода LLM.

TransferEngine устраняет эти барьеры благодаря четырём основным возможностям:

Общие допущения транспорта: вместо ориентации на конкретный NIC, TransferEngine предполагает лишь то, что базовый транспорт является надёжным, но неупорядоченным, то есть представляет собой пересечение гарантий, предоставляемых транспортами RC и SRD.

Синхронизация ImmCounter: новый примитив для уведомления о завершении, который работает без зависимости от порядка сообщений, позволяя одному и тому же коду выполняться как на упорядоченных (RC), так и на неупорядоченных (SRD) транспортах.

Поддержка нескольких NIC: прозрачное управление несколькими NIC на один GPU, автоматическое сегментирование передач для максимизации пропускной способности.

Минимальная поверхность API: небольшой набор операций, достаточный для создания сложных коммуникационных паттернов.

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

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

Каждый экземпляр TransferEngine предоставляет единственный основной адрес для идентификации и обнаружения узлов. Удалённые узлы обмениваются структурами NetAddr, которые кодируют сетевые адреса и могут быть сериализованы для внеполосной связи. Все узлы в системе должны использовать одинаковое количество NIC на каждый GPU. Эта единообразная конфигурация позволяет TransferEngine интеллектуально сегментировать передачи: при инициировании передачи и источник, и получатель точно знают, какие NIC доступны, что позволяет автоматически распределять нагрузку по всем адаптерам.

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.

Для небольших сообщений и контрольных сообщений в стиле RPC, TransferEngine предоставляет submit_send и submit_recvs:

submit_send: отправляет сообщение удалённому узлу, идентифицированному с помощью NetAddr. Сообщение копируется внутренне, поэтому вызывающий может немедленно повторно использовать или освободить буфер.

submit_recvs: размещает пул буферов приёма. При поступлении сообщения вызывается обратный вызов с полученными данными. Буфер автоматически возвращается в пул после завершения обратного вызова.

Для массовых передач данных TransferEngine предоставляет операции односторонней записи с нулевым копированием:

submit_single_write: передаёт непрерывную область от источника к назначению

submit_paged_writes: передаёт несмежные страницы, заданные индексами, шагами и смещениями

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

Движок автоматически сегментирует большие передачи между несколькими NIC для максимизации пропускной способности.

Для операций, направленных на несколько узлов (распространено в маршрутизации MoE), TransferEngine предоставляет оптимизированные примитивы:

add_peer_group: предварительно регистрирует набор адресов узлов для быстрого поиска

submit_scatter: отправляет различные срезы данных каждому узлу в группе

submit_barrier: отправляет уведомление «только немедленное значение» всем узлам (передача нулевого байта)

Эти операции оптимизированы для низкой задержки путём кэширования внутренних структур данных и шаблонизации запросов на выполнение RDMA, что позволяет избежать повторяющихся накладных расходов на настройку.

Ключевая проблема в системах LLM — координация передач RDMA с прогрессом GPU. В идеале передача данных должна начинаться, как только GPU выдает результат, но опрос GPU со стороны CPU неэффективен.

TransferEngine решает эту задачу с помощью UVM-наблюдателей: единых областей виртуальной памяти, которые могут обновляться ядрами GPU (даже внутри графов CUDA) и эффективно опрашиваться CPU с помощью GDRCopy. Когда выделяется UVM-наблюдатель, регистрируется обратный вызов, который вызывается всякий раз, когда GPU обновляет значение.

Это позволяет реализовать такие шаблоны, как послойная передача KvCache: после завершения каждого слоя трансформера GPU увеличивает значение наблюдателя, запуская передачу RDMA KvCache этого слоя, в то время как последующие слои продолжают вычисления.

TransferEngine предоставляет механизмы уведомления о завершении как для стороны инициатора (вызывающего submit_write), так и для стороны получателя (цели операций записи).

Уведомление со стороны инициатора: при отправке операции записи вызывающий может выбрать получение уведомления по завершении через обратный вызов или атомарный флаг. Это уведомление подтверждает, что данные были успешно переданы из исходного буфера.

Уведомление со стороны получателя: ключевой инновацией является ImmCounter, компонент, который отслеживает счетчики, индексируемые 32-битными немедленными числами. Когда операция записи несет немедленное значение, ImmCounter получателя увеличивается атомарно по факту доставки данных. Функция expect_imm_count позволяет получателю зарегистрировать обратный вызов, который срабатывает, когда конкретное немедленное значение достигает целевого количества. Счетчики могут быть прозрачно синхронизированы с GPU через GDRCopy, опрашиваться напрямую или отслеживаться через обратные вызовы.

Этот механизм со стороны получателя обеспечивает надёжную синхронизацию без необходимости упорядочивания сообщений. Он одинаково эффективно работает на NIC ConnectX (которые гарантируют доставку в порядке очереди через транспорт RC) и на NIC EFA (которые доставляют сообщения не по порядку через транспорт SRD).

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

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

Сегментирование внутри DomainGroup является гибким. Передачи могут быть направлены на конкретные NIC по индексу. Одиночная запись может быть разделена. Постраничные передачи, операции scatter и barrier, которые переводятся в несколько записей, могут сегментироваться по всем 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, уменьшая количество звонков дверного звонка к NIC.

Ослабленное упорядочивание: включает 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}
}

Мотивация

TransferEngine

Дизайн API

Регистрация памяти

Двусторонняя передача Send/Recv

Односторонняя запись

Группы узлов и массовые операции

UVM-наблюдатель: сигнализация от GPU к CPU

Уведомление о завершении

Реализация

Аппаратно-специфичные оптимизации

Заключение

Цитата