Cohere выложила код сервера, генерирующего до 1,58 раза быстрее vLLM. Как это работает
На одной видеокарте H100 сервер выдаёт 292 токена в секунду, у vLLM — 185. В полном цикле работы сервера выигрыш меньше: от 1,25 до 1,41 раза.
Cohere выложила на GitHub код сервера для своей модели North Mini Code. Весь шаг генерации в нём выполняет одна программа для видеокарты — «мегаядро». При обработке одного запроса за раз сервер работает в 1,58 раза быстрее популярного движка vLLM1, говорится в блоге компании.
Обычно каждая операция модели запускается на видеокарте отдельной программой — ядром: запустили, дождались, запустили следующую. Мегаядро запускается один раз на весь шаг генерации. Каждый вычислительный блок видеокарты берёт задачи из заранее составленного списка и сам ждёт, пока будут готовы нужные ему данные. Простоев между операциями становится меньше.
При генерации по одному токену видеокарта больше ждёт, чем считает: скорость упирается в то, как быстро веса модели читаются из памяти. У North Mini Code 30 млрд параметров, на каждый токен работают 3,3 млрд, за шаг приходится прочитать около 6,6 ГБ весов. По расчёту Cohere, предел для H1002 — около 470 токенов в секунду. vLLM выдаёт 185 токенов, или 39% предела, а мегаядро — 292, или 62%.
«Мы также обнаружили, что мегаядра писать намного проще, чем гласит их репутация», — отмечают авторы. По их словам, всё мегаядро уместилось в один файл на CUDA. Отдельный компилятор и новая модель программирования не понадобились: обычные ядра для умножения матриц и механизма внимания переписали под общий формат задач.
В полном цикле работы сервер обрабатывает реальные запросы пачками по восемь штук. Здесь генерация ускоряется слабее — в 1,25–1,41 раза. На тестах качества сервер показал результаты, близкие к vLLM. Он совместим с API OpenAI и поддерживает вызов инструментов, так что к нему можно подключить агента-программиста OpenCode3.
Ограничения у сервера такие:
- одновременно обрабатывается не больше восьми запросов;
- пока сервер разбирает текст нового запроса, генерация по остальным стоит на паузе;
- мегаядро используется только для генерации, входной текст обрабатывается обычными ядрами PyTorch.
Идею Cohere взяла у Hazy Research — лаборатории Стэнфордского университета, которая раньше собрала в одно ядро небольшую модель Llama-3.2-1B. По словам Cohere, прежние работы с мегаядрами сводились к компиляторам и демонстрациям, а у неё первый полноценный сервер на такой технологии. Следующим компания выпустит мегаядро для видеокарт RTX Blackwell, в том числе RTX 50-й серии. В планах поддержка форматов FP8 и FP44.
Ранее, 2 сентября, Cohere советовала компаниям отдавать узкие задачи малым моделям и приводила в пример North Mini Code — свою первую модель-агента для программирования.