Если ты когда-нибудь занимался самостоятельным развертыванием языковых моделей на серверах, то знаешь, насколько болезненной бывает борьба за каждый токен в секунду. Движок vLLM давно стал золотым стандартом индустрии: алгоритм PagedAttention, непрерывный батчинг и поддержка распределенного инференса выжимают из видеокарт почти всё возможное. Но когда речь заходит о граничной оптимизации под конкретную архитектуру весов, инженеры неизбежно упираются в потолок.
Традиционно оптимизация вычислительных ядер — это элитарная дисциплина. В мире буквально пара сотен специалистов, способных виртуозно писать на C++, CUDA и низкоуровневом языке Triton от OpenAI. На ручной тюнинг одного оператора уходят недели кропотливых экспериментов с регистрами и разделяемой памятью GPU.
Инженеры Google DeepMind решили проверить дерзкую гипотезу: что если поручить задачу оптимизации инференса открытой модели Gemma 4 флагманской рассуждающей нейросети Gemini 4 Argon? Результат превзошел самые смелые ожидания: модель выступила в роли автономного системного компилятора и выдала восемь pull requests в официальный репозиторий vLLM, разогнав пропускную способность почти в три раза. Давай разберем под капотом, какие конкретно узкие места вскрыл алгоритм.
Спекулятивное декодирование с MTP: как предугадывать токены пачками
Главное проклятие авторегрессионных моделей — последовательная природа генерации. Чтобы сгенерировать фразу из пятидесяти слов, видеокарта обязана сделать пятьдесят последовательных проходов по весам. Даже если вычислительные блоки GPU простаивают, задержка обращения к памяти (Memory Bandwidth) душит производительность.
Архитектура Gemma 4 изначально создавалась с прицелом на Multi-Token Prediction (MTP) — способность модели прогнозировать сразу несколько токенов за один прямой проход. Однако в движке vLLM этот потенциал оставался нераскрытым из-за отсутствия эффективного планировщика.
Gemini 4 Argon спроектировал специализированный модуль спекулятивного драфтера (speculative drafter). Механика работает элегантно: легкий драфтер на лету генерирует цепочку из 4 кандидатов токенов параллельно. Затем основная модель Gemma 4 E2B верифицирует всю пачку за один-единственный шаг прямого прохода. Если токены угаданы верно, они мгновенно попадают в ответ. На структурированном выводе — коде, JSON или логах — точность угадывания зашкаливает, что сразу дает мощнейший буст скорости без потери смысловой точности.
Пять кастомных ядер Triton: слияние операций и борьба с утечками памяти
Но одного спекулятивного декодирования мало: узкие места прятались глубже, на уровне взаимодействия с памятью видеокарты. Argon провел профилирование графа вычислений и сгенерировал пять специализированных ядер на языке Triton.
Меня в этом рефакторинге больше всего впечатлили три инженерных хода:
- Слияние слоев RMSNorm и линейной проекции QKV (Kernel Fusion). В стандартном пайплайне тензор сначала нормализуется, записывается в глобальную память GPU, а затем считывается следующим оператором для умножения матриц. Argon объединил эти шаги в одно слитное ядро: данные обрабатываются прямо в быстрой локальной памяти регистров, исключая лишний круг обращения к видеопамяти.
- Аппаратный оператор RoPE с кэшированием тригонометрии. Поворотное позиционное кодирование (Rotary Position Embedding) требует постоянного вычисления синусов и косинусов. Сгенерированное ядро совмещает кэширование координатных сеток с потоковым умножением.
- Разреженная выборка (sparse gather) в голове GEMV. Это была самая жирная оптимизация: в стандартной реализации vLLM промежуточный буфер временных тензоров отъедал до 1540 МБ видеопамяти. Argon переписал алгоритм на прямое потоковое считывание данных, сократив расход вспомогательной VRAM практически до нуля.
В результате освободившаяся память позволила значительно увеличить размер батча (batch size), не вылетая в ошибку Out of Memory.
Замеры на NVIDIA A100: 108 000 токенов в секунду без потери точности
Цифры говорят сами за себя. Замеры проводились на стандартном серверном ускорителе NVIDIA A100 с объемом памяти 80 ГБ:
- Базовая пропускная способность Gemma 4 E2B на vLLM составляла 37 900 токенов в секунду.
- После применения пула патчей от Gemini 4 Argon пропускная способность взлетела до 108 000 токенов в секунду.
- Итоговый коэффициент ускорения составил 2,85 раза.
При этом верификация логитов показала 100% совпадение с эталонным распределением. Это чистая оптимизация вычислений без квантования весов, урезания контекста или ухудшения интеллекта ответов.
Для серверных провайдеров ускорение почти в три раза означает прямое пропорциональное сокращение серверного парка. Там, где раньше требовался кластер из трех серверов стоимостью в десятки тысяч долларов в месяц, теперь справляется одна стойка.
Что это меняет для индустрии открытых моделей
Этот кейс важен не просто как красивая победа в бенчмарках. Мы присутствуем при рождении нового витка развития софта: рекурсивной оптимизации ИИ силами самого ИИ.
Язык Triton от OpenAI задумывался как способ дать питонистам возможность писать быстрые ядра для видеокарт без погружения в дебри ассемблера CUDA. Но выяснилось, что Triton стал идеальной средой именно для рассуждающих моделей. Нейросеть с легкостью удерживает в контексте карту памяти GPU, учитывает размеры тайлов и проектирует код, не уступающий ручной работе ведущих архитекторов кремниевых гигантов.
Конечно, сгенерированные ядра требуют жесткой валидации: код, оптимизированный под архитектуру Ampere (A100), может вести себя иначе на Hopper (H100) или Blackwell. Кроме того, спекулятивное декодирование раскрывается на полную мощность при детерминированных задачах: если ты выкрутишь температуру генерации на 1.5 для написания стихов, процент принятых токенов упадет, и ускорение будет скромнее. Но фундамент заложен: эпоху ручной рутины низкоуровневых программистов сменяет эра автономных ИИ-компиляторов.