Gemma 4 на AWS Graviton: сравнение vLLM, JAX и PyTorch

Сравнение рантаймов для Gemma 4: как получить честные цифры

Вы думаете, что измеряете производительность модели, но на самом деле сравниваете свои измерительные инструменты? В этой статье мы покажем, как три разных рантайма — vLLM, JAX и PyTorch — обслуживают одну и ту же модель Gemma 4 на одном и том же GPU, и почему только единый бенчмарк даёт честную картину. Вы узнаете, как избежать типичных ошибок и сделать свои тесты достоверными.

Мы развернули три инстанса AWS G5g с одинаковым железом и запустили на них модель google/gemma-4-E2B-it. Каждый использовал свой движок: vLLM, собственную JAX-реализацию и PyTorch с transformers. Чтобы сравнение было корректным, мы написали общий бенчмарк, который измеряет скорость декодирования по времени между токенами на проводе — это единственный показатель, который одинаков для всех OpenAI-совместимых серверов.

Почему три бенчмарка — это не сравнение

Изначально у каждого рантайма был свой тестовый скрипт, который выдавал свою метрику. Три разных числа — это не сравнение, а три отдельные истории. Например, JAX-сервер сообщал скорость декодирования с точностью до одного десятичного знака: 12.8, 12.8, 12.8. Такая повторяемость — не воспроизводимость, а ограничение измерителя. При реальной скорости около 13 токенов в секунду один десятичный знак даёт разрешение всего 0.78%, и вы не можете увидеть эффекты меньше 2%.

Решение заняло два символа — мы добавили второй десятичный знак в формат вывода. После этого первое же измерение показало 12.962, и разница стала видимой.

Единый бенчмарк: что мы измеряем

Мы взяли стандартное определение TPOT (time per output token) из vllm bench serve: (latency - ttft) / (output_len - 1). Это позволяет сравнивать наши числа с опубликованными данными vLLM. Скрипт поддерживает три источника: usage (сервер сам сообщает метрику), stream (измеряем по времени между токенами) и both (используем оба). Режим auto автоматически выбирает подходящий.

python3 sweep.py --base http://<ip>:8000/v1 --out benchmarks/runs/<run> \
    --contexts 64,512,1024,2048,3072,3800 --outputs 32,128 --repeats 3 \
    --decode-source both

Вывод показывает, что сервер JAX выдаёт заниженные значения по сравнению с потоковым измерением (коэффициент 0.98), а PyTorch — ещё ниже (0.95). Это значит, что нельзя просто взять цифру одного рантайма и пересчитать для другого — ошибка составит до 2.6%, что критично при разнице между рантаймами в 24%.

Оборудование и ограничения

Мы использовали инстанс g5g.2xlarge — 8 vCPU, 16 ГиБ RAM, GPU NVIDIA T4G с 15 360 МиБ видеопамяти. Это единственное семейство AWS, где GPU работает на процессоре Graviton (aarch64), что даёт уникальное сочетание архитектуры и вычислительной мощности. Квота на vCPU позволяла запускать не более двух инстансов одновременно, поэтому все эксперименты шли последовательно.

Как мы запускали инстансы

Из-за частого отсутствия свободных мощностей мы написали скрипт, который циклически перебирает зоны доступности с паузой 60 секунд. Вот фрагмент лога:

[12:51:45] round 5 us-east-1c: ❌ AWS InsufficientInstanceCapacity
[12:52:47] us-east-1a: ✅ Launching `i-02e79988a6cbeecbf` (g5g.2xlarge, spot, 1x T4G) in `us-east-1`.

Установка занимала в среднем 113 секунд после запуска (это установка из колеса, а не сборка из исходников). Проверка GPU выполнялась реальным умножением матриц, а не просто чтением конфигурации.

Результаты: скорость декодирования

Вот сводная таблица для одного запроса (batch size = 1):

vLLM лидирует с огромным отрывом. Но почему все так далеки от теоретического максимума в 61.4 ток/с? Мы рассчитали его, исходя из того, что модель читает 4.5 ГиБ весов за шаг декодирования при пропускной способности памяти 277 ГиБ/с. Оказывается, ни один рантайм не упирается в память — все ограничены другими факторами, например, накладными расходами на запуск ядер.

Неожиданный результат: время до первого токена

Мы ожидали, что vLLM будет быстрее в декодировании, но не ожидали 30-кратного преимущества во времени до первого токена (TTFT). При длине входа 3746 токенов vLLM показал 178 мс, тогда как PyTorch — 2339 мс. Это казалось невозможным, ведь префилл требует 14 TFLOP, а T4G выдаёт максимум 20-30 TFLOP/с, что даёт минимум 460 мс.

Разгадка в том, что vLLM использует кэширование префиксов (prefix caching). Метрики показали 94.7% попаданий в кэш — то есть vLLM фактически обработал только 5.3% токенов, потому что мы использовали один и тот же промпт для прогрева и повторных запусков. Остальные рантаймы не имеют такого кэша и каждый раз платят полную стоимость.

Чтобы получить честные цифры, мы добавили случайную строку (нонс) в начало промпта — это ломает кэш, так как общий префикс исчезает. Теперь этот тест включён в набор проверок.

Практические выводы

Если вы хотите получить достоверные результаты при сравнении рантаймов, следуйте этим правилам:

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

Начните с малого: возьмите свой текущий сервер и запустите на нём наш скрипт с режимом stream. Вы увидите, какие цифры реальны, и сможете принимать решения на основе фактов, а не догадок.