AWS показала, как speculative decoding на Trainium2 ускоряет генерацию в vLLM

AWS показала практический способ ускорить и удешевить LLM-инференс на Trainium2 для сценариев, где модель генерирует намного больше токенов, чем получает на вход. Речь о speculative decoding: вместо того чтобы заставлять большую модель последовательно выдавать по одному токену, система подключает небольшую draft-модель, которая быстро предлагает сразу несколько следующих токенов, а основная target-модель проверяет их за один проход. Если предположения совпадают, сервис тратит меньше дорогих последовательных шагов, снижает задержку между токенами и лучше загружает ускоритель.

Это особенно важно для decode-heavy нагрузок — ассистентов для письма, coding agents, генерации отчетов, шаблонных документов и других задач с длинным ответом. В обычной авторегрессионной генерации каждый новый токен считается отдельно, поэтому ускоритель постоянно читает KV-cache из памяти и делает сравнительно мало полезной работы за шаг. Из-за этого инференс часто упирается в пропускную способность памяти, а не в чистые вычисления.

Speculative decoding бьет именно в это узкое место: target-модель реже выполняет последовательные decode-шаги, а верификация пачки токенов делает нагрузку более плотной. Но у подхода есть требования. Draft- и target-модели должны использовать один и тот же токенизатор и словарь, а лучше еще и принадлежать к одному архитектурному семейству, чтобы малая модель чаще угадывала продолжение основной.

Ключевой параметр — число speculative tokens. Если окно слишком маленькое, выигрыш почти незаметен; если слишком большое, ранние отклонения и лишняя проверка съедают эффект. В своем тесте AWS использовала target-модель Qwen3-32B и draft-модель Qwen3-1.

7B, запущенные через vLLM на Trn2-инстансе trn2.48xlarge. Для speculative decoding выбрали fused speculation в NeuronX Distributed Inference, когда обе модели компилируются вместе ради лучшей производительности.

Базовую и speculative-конфигурации развернули в одном кластере Amazon EKS и сохранили все одинаковым: allocation ускорителей, tensor parallelism, длину контекста, batch limits и Neuron-образ. Разница была только в добавлении draft-модели и параметра num_speculative_tokens. Нагрузку на оба сервиса подавали через llmperf, а TTFT, inter-token latency и сквозную задержку отправляли в CloudWatch для сопоставления.

Отдельно AWS проверила и более компактную Qwen3-0.6B, но ее acceptance rate оказался примерно на 60 процентов ниже, и этого хватило, чтобы потерять основную часть выигрыша. В диапазоне от 5 до 15 speculative tokens лучшей точкой в этих тестах стала конфигурация с семью токенами, хотя авторы отдельно подчеркивают, что оптимальное значение сильно зависит от структуры промптов.

Именно структура запроса в итоге определила результат. На предсказуемых сценариях — повторяющийся текст, числовые последовательности, простой код — speculative decoding показал явную пользу. В таких случаях draft-модель часто угадывает то, что и так выдала бы target-модель, поэтому система пропускает существенную часть последовательных шагов.

В тестах inter-token latency опускалась примерно до 15 миллисекунд на токен, а кривая сквозной задержки стабильно шла ниже baseline. На открытых, менее детерминированных запросах картина другая: draft-модель чаще расходится с target-моделью, токены отклоняются, и потенциальный выигрыш исчезает. Для таких промптов inter-token latency держалась около 45 миллисекунд на токен, а speculative и baseline-конфигурации показывали почти одинаковую общую задержку.

TTFT, то есть время до первого токена, при этом почти не изменился, потому что speculative decoding не ускоряет prefill-этап, где модель кодирует входной контекст. Главная выгода появляется позже, в decode-фазе, за счет уменьшения числа дорогих последовательных шагов target-модели. Практический вывод из статьи простой: speculative decoding на Trainium2 — не универсальная кнопка ускорения, а точечная оптимизация под конкретный тип нагрузки.

Если продукт часто генерирует структурированный и предсказуемый вывод — код, извлечение данных, шаблонные отчеты, конфиги, — такой режим может напрямую снизить стоимость выходного токена и поднять пропускную способность без потери качества. Если же у тебя в основном открытый чат со свободной генерацией, эффект может быть минимальным. Поэтому внедрять эту схему стоит только после замеров на собственных промптах, с подбором совместимой draft-модели и окна speculative tokens под реальные сценарии, а не по бенчмаркам в отрыве от продукта.