Cohere ускорила decode North Mini Code до 1,58 раза одним CUDA megakernel

Один persistent kernel заменяет цепочку запусков

Cohere представила serving engine для North Mini Code, в котором весь autoregressive decode step выполняет один persistent CUDA megakernel. Вместо последовательных запусков QKV, attention, MoE и других операций на каждом из 132 Streaming Multiprocessors H100 постоянно живёт один threadblock и читает собственный список небольших задач из global memory.

North Mini Code — модель на 30 миллиардов параметров с 3,3 миллиарда активных параметров на токен. В BF16 каждый шаг должен перенести около 6,6 GB весов и примерно 0,5 GB KV cache при контексте 8K. Cohere оценивает теоретический предел H100 с пропускной способностью HBM 3,35 TB/s примерно в 470 tok/s, тогда как vLLM в этом режиме выдаёт 185 tok/s, или 39% предела.

Megakernel достиг 292 tok/s при batch size 1 — 62% расчётного Speed-of-Light и в 1,58 раза больше baseline. Авторы сообщают, что преимущество сохраняется при других размерах batch и контексте до 256K. Эти цифры относятся к BF16, одной H100 и конкретной архитектуре North Mini Code; они не описывают автоматически другие модели и GPU.

Ускорение приходит из пустот между операциями

При decode небольшие GEMV-подобные операции ограничены прежде всего памятью и задержкой, а не числом FLOPS. Обычный сервер ждёт завершения всей сетки перед запуском следующего kernel. Если 200 tiles распределены по 132 SM, вторая волна занимает только 68 процессоров, а остальные простаивают до общей границы синхронизации.

Megakernel заполняет такие промежутки готовой работой из другой операции. Параллельные transformer layers North Mini Code особенно удобны: attention и MoE читают один нормализованный вход и соединяются только перед residual add и RMSNorm. Освободившийся SM может начать MoE tile, пока другая часть устройства заканчивает attention, если необходимые данные уже готовы.

Два других источника выигрыша — отказ от ложных зависимостей и ранняя загрузка весов. O-proj для конкретной KV-группы стартует после готовности именно её attention output, не ожидая остальные группы. Неизменяемые веса QKV, router или MoE down projection можно заранее переносить из HBM в shared memory, пока ещё считается зависимая activation.

Общий ABI делает большой kernel управляемым

Реализация Cohere находится в одном CUDA-файле и не требует отдельного compiler. Шестнадцать opcodes покрывают dense GEMMs, attention, MoE routing, MoE GEMMs и финальное объединение ветвей. Каждая операция подчиняется одному ABI: 12 warps в трёх warpgroups, фиксированный task descriptor из 32 int32 и общий протокол ожидания и сигнализации.

Один controller warp заранее подгружает descriptors, producer переносит данные, восемь consumer warps выполняют вычисления, storer записывает результат и отмечает готовность. Зависимости представлены счётчиками в global memory. Задача ждёт нужного значения счётчика и после записи результата атомарно увеличивает следующий; memory fences не дают потребителю увидеть сигнал раньше данных.

Эта простота не устраняет риск. Ошибка в числе arrivals способна разблокировать неверную задачу или вызвать deadlock гораздо позже места дефекта. Cohere прямо называет barrier accounting самой трудной частью переноса. Поэтому общий ABI уменьшает архитектурную сложность, но требует явных инвариантов, тестов корректности и проверки каждой новой операции.

Порядок задач оказался отдельным рычагом производительности

Большую часть расписания host строит заранее и раздаёт задачи round-robin по 132 SM. Attention и MoE используют локальное work stealing, потому что длины живых последовательностей и распределение токенов по experts становятся известны только во время шага. Такое сочетание сохраняет управляемый статический порядок и балансирует переменный объём работы.

В ablation при контексте 8K только перестановка волн заметно меняла throughput. При batch size 1 настроенный порядок дал 291 tok/s, interleaved вариант — 282 tok/s, а attention-first — 236 tok/s. Попытка размещать зависимые задачи на одном SM ради cache locality, наоборот, замедлила kernel на 1–2%, поэтому опубликованная версия использует простой round-robin.

Serving engine разделяет ответственность между Python control plane и C++ decode thread. Python принимает запросы, запускает prefill и меняет batch, а C++ выбирает заранее построенное расписание, запускает megakernel и стримит токены. Перед изменением batch Python останавливает decode на границе шага, чтобы GPU не прочитал наполовину обновлённые KV tables или metadata.

End-to-end прирост меньше microbenchmark — и это ожидаемо

В полном серверном тесте при batch size 8 средняя decode throughput была выше vLLM в 1,25–1,41 раза: 1,41× на AIME 2025, 1,25× на GPQA, 1,33× на MMLU-Pro CS, 1,37× на SciCode и 1,28× на LiveCodeBench v6. Прирост меньше 1,58×, потому что prefill остаётся набором PyTorch kernels, временно останавливает decode, а активный batch и MoE routing меняются.

Проверка качества не показала измеримого ухудшения: SciCode дал 38,9% ± 1,6% для megakernel против 38,2% у vLLM, а LiveCodeBench v6 — 70,3% ± 1,1% против 70,3%. Среднее и стандартное отклонение megakernel рассчитаны по семи запускам. Близкие значения поддерживают корректность реализации, но не доказывают идентичность каждого токена.

У текущей версии есть жёсткие ограничения: maximum batch size равен 8, megakernel обслуживает только decode, а новый prefill приостанавливает все декодируемые запросы. В планах — RTX Blackwell, FP8 и FP4, затем datacenter Blackwell и multi-GPU. Поэтому работа Cohere — уже функционирующий сервер, но ещё не универсальная замена любому production inference stack.