Как DeepSeek разогнали обмен данными в своих нейросетях

Когда мы обсуждаем DeepSeek-V3 или R1, обычно всё внимание уходит на количество параметров или качество ответов. Но за кулисами этих моделей стоит огромная инженерная работа по оптимизации «железа». Оказалось, что стандартных библиотек для связи между видеокартами ребятам не хватило, и они написали свою — DeepEP.

Зачем нам ещё одна библиотека связи

Если вы когда-нибудь пробовали обучать модели с архитектурой Mixture-of-Experts (MoE), то знаете главную «боль»: данные постоянно летают между экспертами, которые разбросаны по разным GPU и серверам. Этот процесс называют MoE dispatch (рассылка токенов нужным экспертам) и combine (сбор результатов обратно).

Проблема в том, что стандартные операции вроде All-to-All часто становятся узким местом. Пока GPU ждут данные по сети, они простаивают. DeepEP — это низкоуровневая библиотека на CUDA, которая делает пересылку данных максимально быстрой именно для MoE-сценариев. Она заточена под специфическую топологию современных кластеров, где внутри сервера у нас быстрый NVLink, а между серверами — RDMA (InfiniBand или RoCE).

Что умеет DeepEP

Разработчики выделили два основных режима работы, потому что задачи обучения и инференса сильно отличаются по своим требованиям.

Обычные ядра для обучения и префиллинга

В этих задачах нам важна пропускная способность. DeepEP умеет эффективно перекидывать данные из домена NVLink в домен RDMA. Интересная фишка: библиотека позволяет контролировать количество SM (потоковых мультипроцессоров), которые заняты коммуникацией. Это помогает найти баланс между скоростью передачи и вычислительной мощностью, оставшейся для самой нейросети.

Low-latency ядра для генерации текста

Когда модель уже «отвечает» пользователю (этап декодинга), задержки становятся критичными. Здесь DeepEP использует чистый RDMA, чтобы минимизировать любые накладные расходы.

Самое крутое тут — механизм перекрытия (overlapping). Библиотека использует хуки, которые позволяют передавать данные в фоне, пока GPU занят вычислениями (например, Attention), причём это вообще не отъедает ресурсы SM.

Технические хитрости и «недокументированные» возможности

В коде DeepEP есть пара моментов, которые заставят любого системного программиста поднять бровь. Разработчики честно признаются, что ради экстремальной производительности они использовали неочевидное поведение инструкций PTX на архитектуре Hopper (H100/H800).

Например, они используют кэш L1 как некогерентный, читая волатильные данные через ld.global.nc. По документации это может привести к ошибкам, но на практике на Hopper это работает быстрее и корректно, так как L1 и некогерентный кэш там объединены. Если у вас что-то пойдёт не так на других картах, это поведение можно отключить флагом DISABLE_AGGRESSIVE_PTX_INSTRS.

Как это выглядит в коде

Интерфейс библиотеки ориентирован на PyTorch. Сначала нужно инициализировать буфер, причём для режима низкой задержки памяти потребуется заметно больше.

from deep_ep import Buffer

# Настраиваем количество SM для коммуникации
Buffer.set_num_sms(24)

# Инициализируем буфер
buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes)

# Пример dispatch-операции
recv_x, recv_topk_idx, _, _, handle, event = buffer.dispatch(
    x, topk_idx=topk_idx, topk_weights=topk_weights,
    num_tokens_per_rank=num_tokens_per_rank,
    async_finish=True
)

Для инференса всё ещё интереснее — можно использовать хуки для того самого «бесплатного» перекрытия вычислений и передачи данных:

recv_hidden_states, _, handle, event, hook = buffer.low_latency_dispatch(
    hidden_states, topk_idx, num_max_tokens, num_experts,
    return_recv_hook=True
)

# Вызываем hook(), когда данные реально понадобятся
# До этого момента передача идет в фоне, не занимая SM
hook()

Производительность на практике

Цифры в README выглядят солидно. На тесте с H800 и сетью 400 Гбит/с они получают около 150 ГБ/с внутри узла через NVLink и стабильные 40-50 ГБ/с между узлами через RDMA. При этом задержки в low-latency режиме для небольших батчей (128 токенов) укладываются в 77–190 микросекунд даже при масштабировании на 128-256 экспертов.

Кстати, в апреле 2025 года инженеры из Tencent закинули пул-реквест, который поднял производительность ещё на 30%. Так что проект живой и активно оптимизируется сообществом.

Кому стоит попробовать

DeepEP — это не инструмент для домашнего обучения классификатора котиков. Вам он пригодится, если:

1. Вы строите или эксплуатируете MoE-модели уровня DeepSeek.
2. У вас есть кластер с NVLink и RDMA (InfiniBand/RoCE).
3. Вы упёрлись в производительность стандартных коллективных операций NCCL.

Библиотека требует CUDA 12.3+ для архитектур SM90 (Hopper) и NVSHMEM в качестве зависимости. Если вы работаете с архитектурой Ampere (A100), поддержка тоже есть, но только для внутриузловых коммуникаций.

В целом, это отличный пример того, как глубокая оптимизация под конкретное железо позволяет выжимать максимум из дорогих вычислительных ресурсов. Если вы занимаетесь распределённым обучением, заглянуть в исходники DeepEP стоит хотя бы ради того, чтобы посмотреть на хитрые трюки с PTX и организацией очередей в памяти GPU.