Как ускорить дешёвый декодинг LLM на AWS Trainium: практический разбор speculative decoding с Qwen3 и vLLM — VogueTech

Что нового

AWS показала, как выжать больше производительности из своих чипов Trainium2 для генеративных моделей с длинными ответами за счёт speculative decoding в связке с vLLM и Kubernetes.

Ключевые факты из экспериментов с Qwen3:

Используются пары моделей Qwen3-32B (таргет) + Qwen3-1.7B (draft) на Trainium2 (инстансы trn2.48xlarge).
Спекулятивный декодинг ускоряет генерацию токенов для decode‑heavy задач до 3 раз по сравнению с обычным автодекодингом, при сохранении качества вывода.
На структурированных промптах (повторы, числовые последовательности, простой код) межтокенная задержка падает примерно до 15 мс на токен.
Для открытых промптов (например, «I believe the meaning of life is») межтокенная задержка остаётся около 45 мс на токен, сопоставимо с базовой конфигурацией без speculative decoding.
Время до первого токена (TTFT) практически не меняется — speculative decoding ускоряет только стадию декодинга, а не префилл.
Лучший баланс для тестового набора промптов: Qwen3-1.7B в роли draft‑модели и окно из 7 speculative токенов (num_speculative_tokens=7).
Более маленькая draft‑модель Qwen3-0.6B работает быстрее, но даёт примерно на 60% меньший процент принятых токенов, и итоговый выигрыш по скорости пропадает.
Реализация использует NeuronX Distributed Inference (NxDI) с режимом fused speculation (enable_fused_speculation=true), где draft‑ и target‑модели компилируются вместе под Trainium.

AWS выложила полный набор артефактов:

Манифесты Kubernetes для двух сервисов vLLM (базовый и speculative) на Trn2.
Конфиги vLLM для включения fused speculative decoding.
Скрипты llmperf для нагрузочного тестирования и сбора метрик (TTFT, inter‑token latency, end‑to‑end latency, throughput).
Примеры настройки S3 CSI Driver для монтирования чекпоинтов и скомпилированных артефактов Neuron.
Рекомендации по настройке Neuron DRA, tensor parallelism и раскладки NeuronCore.

Как это работает

Базовая идея speculative decoding

Спекулятивный декодинг ускоряет авторегрессию за счёт двух моделей, работающих в паре:

Draft‑модель (меньшая, быстрее) по текущему контексту предлагает сразу n следующих токенов.
Target‑модель (большая, основная) одним форвардом проверяет эти n токенов и либо принимает их, либо частично/полностью отклоняет.

Критические требования:

Draft‑ и target‑модели должны использовать один и тот же токенайзер и словарь.
Лучше всего работают модели одного архитектурного семейства (как Qwen3-1.7B + Qwen3-32B), потому что их прогнозы следующего токена чаще совпадают.

Основной параметр, которым вы управляете:

num_speculative_tokens — сколько токенов draft‑модель предлагает за один раз.

Если target‑модель принимает предложенные токены, система «перепрыгивает» несколько последовательных шагов декодинга и экономит время и память.

Откуда берётся ускорение

Ускорение складывается из двух эффектов:

Меньше decode‑шагов у target‑модели.
- В обычном автодекодинге каждый новый токен — это отдельный проход с чтением всего KV‑кеша из памяти.
- KV‑кеш хранит ключи и значения для всех прошлых токенов; каждый decode‑шаг читает его целиком, и вычисления становятся ограничены пропускной способностью памяти, а не вычислительными блоками.
- При speculative decoding часть этих шагов пропадает, потому что несколько токенов принимаются «оптом».
Более плотная загрузка железа на этапе верификации.
- В стандартном режиме один decode‑шаг даёт один токен, а матричные умножения в ускорителе работают ради этого одного токена — вычислительные блоки простаивают.
- В режиме верификации target‑модель сразу обрабатывает n токенов, что делает задачу более compute‑dense: меньше накладных расходов на память, больше полезной математики за один запуск.

Если num_speculative_tokens слишком маленький — ускорение ограничено, потому что вы всё ещё делаете почти столько же шагов декодинга.

Если num_speculative_tokens слишком большой — растёт число ранних отклонений токенов, и вы тратите лишние ресурсы на draft‑модель и на частые верификации.

Выбор draft‑модели и окна спекуляции

AWS протестировала две связки draft‑моделей:

Qwen3-0.6B — быстрее, но примерно на 60% меньше принятых токенов, итоговая скорость хуже.
Qwen3-1.7B — чуть тяжелее, но даёт гораздо более высокий процент совпадений с Qwen3-32B.

По num_speculative_tokens проверяли значения от 5 до 15:

Окно 5 токенов даёт небольшой прирост скорости.
Окно 15 токенов увеличивает количество отклонений и в сумме ухудшает производительность.
Лучший баланс для тестовых промптов: 7 speculative токенов с Qwen3-1.7B.

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

Что даёт NeuronX Distributed Inference (NxDI)

AWS Neuron — это SDK для чипов Trainium и Inferentia. В нём библиотека NeuronX Distributed Inference (NxDI) отвечает за масштабируемый inference LLM.

NxDI поддерживает четыре режима speculative decoding на Trainium:

Vanilla speculative decoding
- Draft‑ и target‑модели компилируются отдельно.
- Самый простой способ начать.
Fused speculation
- Draft‑ и target‑модели компилируются совместно.
- Меньше накладных расходов, выше производительность.
- Именно этот режим использует AWS в описанном сетапе (enable_fused_speculation=true).
EAGLE speculation
- Draft‑модель использует скрытые состояния target‑модели.
- Это повышает процент принятых токенов.
Medusa speculation
- Несколько маленьких prediction‑голов работают параллельно и предлагают токены.
- Это уменьшает накладные расходы на draft‑модель.

В описанном эксперименте используется связка Qwen3-1.7B (draft) + Qwen3-32B (target) в режиме fused speculation.

Инфраструктурный сетап

AWS развернула два сервиса vLLM на одном кластере Amazon EKS с Trainium2:

Базовый сервис qwen-vllm — Qwen3-32B с обычным декодингом.
Спекулятивный сервис qwen-sd-vllm — тот же Qwen3-32B как target + Qwen3-1.7B как draft, num_speculative_tokens=7.

Общее для обоих сервисов:

Инстансы Trn2 (trn2.48xlarge).
Одинаковое количество NeuronCore и та же схема tensor parallelism.
Одинаковая максимальная длина последовательности и лимиты батчинга.
Одна и та же Neuron DLC‑картинка.

Единственное отличие — добавление draft‑модели и включение speculative decoding во втором сервисе.

Для нагрузочного тестирования использовался llmperf:

Генерация идентичных паттернов запросов к обоим endpoint‑ам.
Метрики TTFT, inter‑token latency и end‑to‑end latency отправлялись в Amazon CloudWatch.
Telemetry инфраструктуры собиралась через CloudWatch Container Insights.

Что это значит для вас

Для каких задач это реально полезно

Спекулятивный декодинг особенно интересен, если вы строите:

AI‑ассистенты для письма (копирайтинг, длинные письма, отчёты).
Код‑агенты и автодополнение кода.
Генераторы конфигураций, шаблонных документов, SQL‑запросов.
Системы структурированного извлечения данных, которые выдают предсказуемый формат.

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

В этих сценариях speculative decoding на Trainium даёт вам:

До 3x ускорения генерации токенов для decode‑heavy нагрузок.
Снижение стоимости за токен вывода, потому что вы лучше загружаете чипы и делаете меньше decode‑шагов target‑модели.
Рост пропускной способности без потери качества текста, если правильно подобрать draft‑модель и окно спекуляции.

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

Повторяющиеся шаблоны текста.
Числовые последовательности.
Простой, детерминированный код.

Там draft‑модель почти всегда угадывает, что сгенерирует target‑модель, и система экономит максимум ресурсов.

Когда пользы не будет

Спекулятивный декодинг не универсален.

Сценарии, где он почти не помогает:

Открытые, творческие промпты: эссе, рассуждения, креативный текст.
Любые запросы, где модель может пойти по множеству разных траекторий, даже при temperature=0.0.

В этих случаях:

Draft‑модель часто расходится с target‑моделью.
Процент принятых токенов падает.
Вы тратите ресурсы на draft‑модель и верификацию, но не выигрываете в скорости.

AWS прямо показывает это в результатах:

Для открытых промптов межтокенная задержка остаётся ~45 мс на токен и для базового, и для speculative сервиса.
Кривые end‑to‑end latency почти совпадают.

Ещё один важный момент:

TTFT (Time To First Token) не меняется, потому что speculative decoding никак не ускоряет префилл.
Если ваша задача чувствительна именно к первому токену (например, чат с короткими ответами), speculative decoding мало поможет.

Доступность и ограничения

Всё описанное относится к AWS Trainium2 (Trn2) и стеку AWS Neuron + vLLM + Amazon EKS.
Для развёртывания вам нужен доступ к AWS, настройка Kubernetes‑кластера и работа с Neuron SDK.
В России доступ к AWS часто требует VPN и юридической оценки рисков использования зарубежного облака.

Если у вас уже есть инфраструктура на AWS и вы планируете крупные decode‑heavy нагрузки, есть смысл протестировать speculative decoding на части трафика и померить фактическую экономию.

Место на рынке

Спекулятивный декодинг как подход уже используется в разных стеках, но в этом материале AWS фокусируется на Trainium2 + NxDI + vLLM и модели Qwen3.

Что можно сказать по фактам:

AWS заявляет до 3x ускорения генерации токенов для decode‑heavy задач на Trainium2.
Для структурированных промптов межтокенная задержка падает до ~15 мс против ~45 мс для открытых промптов.
Для открытых промптов выигрыш по времени почти исчезает.

Прямых числовых сравнений с GPT‑4o, Claude 3.5, Llama 3.1 или другими проприетарными стеками в материале нет. Поэтому корректно говорить только о сравнении внутри Trainium‑экосистемы:

Спекулятивный декодинг даёт ощутимый прирост относительно того же Qwen3-32B на тех же Trn2, но без speculative режима.
Стоимость за токен вывода на Trainium2 с speculative decoding ниже, чем на том же железе без него — за счёт лучшей загрузки NeuronCore и меньшего числа decode‑шагов.

Если вы уже используете Trainium2 для inference больших моделей, speculative decoding — это способ снизить стоимость и задержки без смены железа и модели.

Как запустить

Исходный материал даёт архитектуру и перечисляет артефакты, но не приводит полный листинг команд. Ниже — структурированный план, который повторяет их сетап и укажет, где искать готовые файлы.

1. Подготовить кластер Amazon EKS с Trn2

Создайте кластер Amazon EKS.
Добавьте node group с инстансами trn2.48xlarge.
Установите AWS Neuron и драйверы на узлах (см. документацию AWS Neuron).

2. Развернуть два сервиса vLLM

В репозитории AWS Neuron EKS samples вы найдёте:

Манифесты Kubernetes для:
- Базового сервиса vLLM с Qwen3-32B (qwen-vllm).
- Сервиса vLLM со speculative decoding (qwen-sd-vllm).
Примеры конфигов vLLM для включения fused speculation.

Общая логика:

Смонтировать чекпоинты моделей Qwen3-32B и Qwen3-1.7B через S3 CSI Driver.
Настроить Neuron DRA и tensor parallelism, чтобы веса Qwen3-32B корректно распределились по NeuronCore.
Для speculative сервиса включить флаг enable_fused_speculation=true и указать:
- Target‑модель: Qwen3-32B.
- Draft‑модель: Qwen3-1.7B.
- num_speculative_tokens=7.

3. Настроить llmperf для бенчмарка

В том же репозитории есть:

Манифест пода qwen-llmperf-pod.yaml.
Скрипты для запуска llmperf с нужными параметрами.
Набор промптов: от структурированных (повторы, числа, простой код) до открытых текстов.

Типичный сценарий:

Запустить под с llmperf.
Настроить его на одновременную отправку запросов к обоим endpoint‑ам:
- Базовый vLLM.
- vLLM со speculative decoding.
Использовать фиксированные длины входа и выхода и temperature=0.0, чтобы нагрузить именно детерминированный декодинг.
Включить логирование и экспорт метрик:
- TTFT (prefill).
- Inter‑token latency (decode).
- End‑to‑end latency.
- Throughput.

4. Мониторинг и анализ

Подключите CloudWatch Container Insights для метрик инфраструктуры.
Отправляйте пользовательские метрики llmperf в CloudWatch dashboards.
Сравните кривые латентности для двух сервисов по разным типам промптов.

5. Тюнинг под свои задачи

После базового запуска имеет смысл:

Подставить свои реальные промпты и длины ответов.
Попробовать разные draft‑модели из одной семьи (например, другие размеры Qwen3).
Прогнать num_speculative_tokens в диапазоне 5–15 и посмотреть, где максимум выигрыша.
Отслеживать процент принятых токенов: если он низкий, уменьшайте окно или улучшайте draft‑модель.

Выводы для практики

Если вы уже на AWS и планируете крупные decode‑heavy нагрузки (ассистенты, код‑генерация, отчёты), имеет смысл закладывать speculative decoding на Trainium2 в архитектуру сразу.
Реальный прирост зависит от структуры промптов: чем более предсказуем формат ответа, тем ближе вы к заявленным до 3x ускорения.
Для креативных и сильно вариативных запросов speculative decoding даёт мало пользы — там лучше оптимизировать другие части пайплайна (модель, батчинг, маршрутизацию).
TTFT не ускоряется, поэтому для сценариев с короткими ответами и жёсткими требованиями к первому токену нужно смотреть на другие оптимизации.

Если ваша команда готова работать с AWS Neuron и Kubernetes, вы можете воспроизвести сетап из репозитория AWS Neuron EKS samples и получить сопоставимые метрики на своих данных.