Команда vLLM разогнала Kimi K3 на Google TPU на 57% быстрее, чем на GPU NVIDIA

Стартап Inferact, который основали создатели и ключевые мейнтейнеры открытого движка vLLM, показал неожиданный результат: на 16 тензорных ускорителях Google TPU v7 модель Kimi K3 выдаёт 709 токенов в секунду, тогда как на 16 GPU NVIDIA GB200 — только 452. Разница 57% в пользу TPU, причём по паспорту железо у NVIDIA мощнее. Секрет не в чипах, а в программной прослойке: команда написала «мегакернел», который собирает сотни мелких операций вывода в одну большую программу. Код выложен в открытый доступ.

Что именно сравнивали

Inferact поставила TPU и GPU в одинаковые условия: 16 ускорителей с каждой стороны, одна и та же модель Kimi K3, один и тот же движок vLLM. Менялись только чипы и низкоуровневая реализация ядер. Со включённым спекулятивным декодированием TPU выдал 709 токенов в секунду против 452 у GB200.

Если отключить спекулятивное декодирование и смотреть «голую» скорость, разрыв не исчезает. При размере батча 1 TPU делает 249 токенов в секунду против 127 у GB200 — почти вдвое больше. При батче 8 — 865 против 636. На другой модели, Qwen 3.8 27B, разрыв ещё заметнее: 4 чипа TPU дают 1515 токенов в секунду, а 4 GPU — 695.

Почему более «слабый» чип оказался быстрее

По паспорту TPU v7 Ironwood и GB200 близки: пиковая производительность в BF16 — 2,31 против 2,5 PFLOPS, в FP8 — 4,61 против 5 PFLOPS. А по пропускной способности памяти TPU даже проигрывает: 7380 ГБ/с против 8000 ГБ/с. То есть у NVIDIA память быстрее, а вывод — медленнее. Значит, дело не в железе, а в том, как оно используется.

Вывод языковой модели упирается не в вычисления, а в перекладывание данных. Чтобы сгенерировать один токен, нужно прогнать через чип все веса модели: вытащить их из памяти HBM во внутренний кэш, посчитать, а на следующем токене повторить всё заново. Обычно этот процесс разбит на сотни отдельных маленьких программ — ядер. Пока одно ядро заканчивается, а следующее ещё не запустилось, память простаивает. Эти паузы складываются в заметную потерю.

Мегакернел: одна программа вместо сотни

Идея Inferact — убрать границы между ядрами. Вся модель Kimi K3, а это 92 слоя смеси экспертов, умещается в одну программу на языке Pallas — это низкоуровневый язык для TPU, аналог CUDA у NVIDIA. Раз границ нет, веса следующего слоя можно начинать подгружать заранее, пока текущий слой ещё считает. Пока слой N занят своими экспертами, веса внимания слоя N+1 уже едут из памяти на чип. Память почти не простаивает.

Такой трюк хорошо ложится на архитектуру TPU. У каждого тензорного ядра TPU v7 есть 64 МиБ быстрой памяти на кристалле, и программист сам решает, что там лежит и когда это освободить. У GPU память устроена иначе: около 38 МиБ на всю карту, разбитые между 152 блоками и управляемые железом, — заранее разложить веса туда негде.

Побочный, но приятный эффект — скорость сборки. Обычный компилятор TPU собирает большую модель больше 30 минут, а мегакернел Inferact компилируется меньше чем за 90 секунд. Для инженера это разница между «правку проверишь завтра» и «правку проверишь через минуту».

Спекулятивное декодирование и точность

Второй слой ускорения — спекулятивное декодирование по схеме DSpark, которую предложила команда DeepSeek. Маленькая модель-черновик быстро предлагает несколько вариантов следующего токена, а большая модель проверяет их пачкой: угаданные принимаются сразу, ошибочные отбрасываются. На TPU средняя длина принятой цепочки дошла до 6 токенов, а один шаг декодирования занимает около 8,5 миллисекунды.

Важно, что за скорость не пришлось платить качеством. На жадном декодировании с максимальным «размышлением» мегакернел показал 94,4% на GPQA-Diamond и 97,2% на GSM8K — ровно те же цифры, что и на GPU.

Кто это сделал

Inferact — компания выходцев из команды vLLM, самого популярного открытого движка для вывода нейросетей. CEO Саймон Мо — один из мейнтейнеров vLLM, соучредитель Вусук Квон — автор проекта и алгоритма PagedAttention, главный научный сотрудник Ю Кайчао пришёл из Университета Цинхуа. В этом году компания подняла 150 миллионов долларов посевных инвестиций при оценке 800 миллионов. Работа над мегакернелом — совместный проект с Google Cloud, и весь код открыт.

Что это значит на практике

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

Оговорки тоже честные. Мегакернел пока заточен под конкретную архитектуру Kimi K3 и конкретную конфигурацию из 16 чипов — под другую модель или другое число ускорителей его придётся переписывать. Это первый шаг, а не готовое универсальное решение.

Частые вопросы

Что такое мегакернел?

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

TPU теперь лучше GPU для нейросетей?

Не так однозначно. Здесь сравнивали узкий сценарий — вывод одной модели на небольшом батче. В других задачах, особенно на обучении, расклад может быть иным. Но результат показывает, что программная оптимизация способна перевесить разницу в железе.

Можно ли попробовать Kimi K3?

Да, веса модели открыты. Поговорить с ней можно в чате NeuralSpace, а агентные сценарии попробовать в разделе «Код».