Спекулятивное декодирование ускоряет генерацию Больших Языковых Моделей (LLM), используя быструю и компактную модель черновика для создания кандидатов завершения, которые проверяются более крупной целевой моделью.

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

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

Спекулятивное декодирование использует структуру естественных языков и авторегрессионную природу трансформеров для ускорения генерации токенов. Хотя более крупные модели, такие как Llama-70B, содержат больше знаний, чем более мелкие, такие как Llama-1B, в некоторых более простых задачах они работают схоже. Это пересечение указывает на то, что определенные последовательности лучше генерируются более дешёвыми моделями, оставляя сложные проблемы для более крупных моделей. Проблема заключается в том, чтобы определить, какие завершения лучше, и равна ли по качеству генерация меньшей модели генерации более крупной.

К счастью, LLM являются авторегрессирующими трансформерами: получив последовательность токенов, они выдают распределение вероятностей для следующего токена. Кроме того, логиты, полученные из промежуточных признаков, связанных с токенами в последовательности ввода, также указывают на вероятность того, что модель выдаст именно эти токены. Это свойство позволяет использовать предположения: если последовательность токенов создается более мелкой моделью, начиная с префикса ввода, она может быть обработана более крупной моделью, чтобы проверить, насколько хорошо она согласуется с целевой моделью. Каждый префикс кандидатов оценивается с вероятностью, и выбирается самый длинный выше порога принятия. В довершение ко всему целевая модель предоставляет последующий токен бесплатно: если черновая модель генерирует n токенов, можно выдать до n + 1 за один шаг.

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

1. Заполнение: как целевые, так и черновые модели должны быть запущены на входной последовательности, чтобы заполнить KV-кэш. В некоторых схемах, таких как Medusa, используют более простые плотные слои для прогнозирования, в этом посте мы сосредоточимся на черновиках на основе трансформеров, которым требуются собственные KV-кэши.

2. Генерация черновика: черновая модель повторяется, чтобы произвести определенное количество фиксированных токенов. Черновая последовательность может быть линейной, или модель может исследовать структуру подобную дереву до определенной глубины (EAGLE, Medusa). Здесь мы сосредоточимся на линейных последовательностях.

3. Принятие: целевая модель запускается на черновой последовательности, создавая логиты, соответствующие каждому черновому токену. Определяется длина самой длинной приемлемой последовательности.

4. Генерация цели: поскольку целевой сгенерировал логиты, в несоответствующей позиции или в конце последовательности логиты соответствуют последующему токену. Эти логиты могут быть выбраны для предоставления нового токена из цели, завершая последовательность.

Существуют различные методы для реализации спекулятивного декодирования. В этом посте мы сосредоточимся на схемах, которые мы использовали для ускорения моделей Sonar, используя внутреннюю модель 1B, а также механизмы прогнозирования, которые мы разрабатываем для ускорения моделей масштаба DeepSeek.

Цель-Черновик

Спекулятивное декодирование может быть достигнуто путем связывания существующей малой LLM в качестве черновой модели с целевой моделью для генерации кандидатных последовательностей. В производственной среде мы ускорили Sonar с использованием модели Llama-1B, настроенной на тот же набор данных, что и целевая. Хотя этот подход не требовал обучения черновика с нуля, малая модель все же использует значительный объем емкости KV-кэша и вводит незначительное заполнитель, увеличивая TTFT.

В этой схеме декодер делает предположения только для пакетов, в которых выполняется только декодирование, создавая токены через стандартную выборку во время заполнения или на смешанных заполнительной-декодирующей стадиях. На этапе заполнения целевые логиты сразу же выбираются, чтобы также заполнить только что сгенерированный токен в KV-кэше черновика. Черновик еще не выбран, но логиты, которые он производит, переносятся на этап декодирования.

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

EAGLE

EAGLE — это схема спекулятивного декодирования, которая исследует несколько черновых последовательностей, сгенерированных через деревообразный обход вероятных черновых токенов. Фиксированное дерево (EAGLE) или динамично-формируемое дерево (EAGLE-2) исследуется с использованием последовательных выполнений черновых токенов, рассматривая кандидаты Top-K в каждом узле вместо следования за токеном с самым высоким рейтингом в линейной последовательности. Затем последовательности оцениваются, и самая длинная подходящая выбирается для продолжения, также добавляя дополнительный токен из цели.

Чтобы добиться более точного прогнозирования, черновая модель EAGLE предсказывает не только на основе токенов, но и используя целевые признаки (скрытые состояния последнего слоя) целевой модели. Недостатком EAGLE является необходимость в обучении индивидуальных, малых черновых моделей, которые достаточно точны для генерации подходящих кандидатов с малой задержкой. Обычно черновая модель представляет собой один слой трансформера, идентичный слою декодера оригинальной модели, который плотно связан с целью, привязываясь к её встраиваниям и проекциям lm_head. Поскольку это требует менее ёмкости KV-кэша, EAGLE имеет меньший объём памяти.

Чтобы проверить древоподобные последовательности в целевой модели, необходимы специальные маски внимания. К сожалению, использование специальной маски внимания для всей последовательности значительно замедляет внимание для реалистичных длин ввода (до 50%), нивелируя некоторые ускорения, достигаемые за счёт предположений. По этой причине мы еще не развернули полное исследование деревьев в продакшене, сосредоточившись на специальном случае однотокенового прогнозирования с помощью схем, подобных MTP, представленных в техническом отчете DeepSeek-V3.

MTP

Эта схема схожа с декодированием черновик-цель, за исключением использования скрытых состояний вместе с токенами для прогнозирования. Требуется выполнить немного больше работы как на этапах заполнителя, так и на этапах декодирования по сравнению с обычными предположениями черновик-цель. Черновая модель использует как токены, так и скрытые состояния: токен t_{i+1} выбирается из логитов L_i, соответствующих токену t_i, которые в свою очередь выводятся из скрытых состояний H_i. Следовательно, буферы входных токенов должны сдвигаться на один шаг влево относительно векторов скрытых состояний, выводимых целью. На изображении ниже отмечены соответствия, используемые для обучения, а также сдвиг в процессе вывода.

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

Обучение голов MTP

Чтобы извлечь выгоду от MTP, мы создали инфраструктуру, необходимую для обучения голов MTP, прикрепленных к нашим настроенным моделям на наборах данных Perplexity, работающей на одном узле с 8 устройствами H100. За один день мы можем создать головы для моделей от Llama-1B до Llama-70B и DeepSeek V2-Lite. Для больших моделей мы полагаемся на головы MTP, созданные в процессе дополнительной настройки.

Цель обучения MTP — сопоставить скрытые состояния черновика и логиты, экстраполированные из скрытых состояний цели, с следующими токеновыми логитами и скрытыми состояниями цели. Поскольку вывод скрытых состояний дорогой, мы предварительно вычисляем их с использованием нашей оптимизированной для вывода реализации целевой модели, чтобы использовать их во время обучения. Однако для проверки реализации вывода MTP и обеспечения того, что численные различия из-за квантования или оптимизаций не мешают результатам, для оценки потерь проверки и точности мы полностью используем реализации вывода как целевой, так и черновой моделей.

При масштабировании с набора данных ShareGPT, использованного в оригинальной статье, до больших выборок, мы заметили, что архитектура голов MTP, изложенных и реализованных в статье EAGLE, не удалось тренировать для моделей размера 70B. В отличие от ShareGPT, который содержал большее количество коротких последовательностей, мы обучаем на немного меньшем количестве значительно более длинных запросов. Поскольку оригинальные головы EAGLE слегка разошлись в структуре от типичного трансформера, мы повторно ввели некоторые слои RMS нормализации, которые были исключены. Мы обнаружили, что это не только позволило обучению сходиться, но и повысило точность голов на несколько процентных пунктов.

Не только слои нормализации способствуют тренировке, повторное введение норм также математически интуитивно. Головки MTP повторно используют встраивания и проекции логитов целевой модели, поскольку они могут быть значительными по размеру (около 2 ГБ для Llama 70B). Во время обучения они заморожены, и ожидание заключается в том, что слой MTP учится встраивать предсказания в то же векторное пространство, что и слой проекций оригинальной модели во время обучения. Отбросив нормы, одиночная MLP ожидается, что она обучится выполнять ту же функцию, что и MLP, за которой следует норма, что мешает сопоставлению скрытых состояний черновика и целевой моделей.

Вывод с использованием спекулятивного декодирования

В движке вывода, для генерации токенов для входных последовательностей, их сначала необходимо сгруппировать в достаточно большие партии, затем должны быть выделены страницы в KV-кэше для следующих токенов. Входные токены и информация о странице KV затем упаковываются в буфер, транслируемый всем параллельным ранкам, работающим с моделью. Наконец, метаданные копируются в память GPU, и модель выполняется для создания логитов, из которых выбирается следующий токен.

В отличие от некоторых реализаций, которые слабо связывают черновой и целевой сервер вывода через обертку, которая управляет запросами между ними, наши пары черновик-цель плотно связаны и шаг за шагом генерируют в унисон. Планирование пакетов и выделение страниц KV объединены между моделями для всех форм спекулятивного декодирования: это объединяет логику, которая связывает модель с вышестоящим сервером вывода, так как они все имеют один и тот же интерфейс.

Время выполнения вывода в Perplexity формируется вокруг FlashInfer, который определяет метаданные, которые необходимо создать для настройки и планирования ядра внимания. Учитывая некоторые входные последовательности, формирующие пакет, для заполнителя, декодирования или проверки, работа на стороне процессора должна быть выполнена для выделения промежуточных буферов и заполнения некоторых постоянных буферов, используемых в внимании. Эта работа добавляется к стоимости планирования пакета и выделению страниц KV, которые также создают задержки, которые необходимо скрыть, чтобы максимально использовать GPU.

Хотя мы полностью параллелизовали работу на стороне процессора и GPU для вывода без предположений, мы обнаружили, что баланс CPU-GPU для спекулятивного декодирования более сложен. Основная проблема возникает из-за того, что количество принятых токенов определяет длину последовательности для последующего выполнения, вводя труднодостигаемый момент синхронизации GPU-to-CPU. Мы экспериментировали с различными схемами планирования, чтобы лучше скрыть задержку работы CPU.

Расписание Черновик-Цель

Несмотря на то, что по размеру меньше целевой модели, когда вся LLM используется в качестве черновика, она все же вводит значительную задержку на GPU, предоставляя некоторое пространство для скрытия затратных операций на CPU. Поскольку меньшие модели не выигрывают от параллелизма тензоров, существует несоответствие между количеством ранков, на которые разделены цель и черновик. В нашей реализации черновая модель работает только на ведущем ранке группы TP.

Как указано ранее, шаг декодирования переносит логиты в следующую операцию. Это позволяет нам наложить одно выполнение черновой модели на работу на стороне CPU по планированию пакета. После того как пакет собран, повторные вызовы к выборщику и черновику создают черновые токены. Параллельно составляется пакет для проверки для целевой модели и синхронизируется с параллельными работниками. Целевые логиты проверяются и отбираются для определения длины принятых последовательностей. На этом этапе необходима синхронизация GPU-to-CPU, чтобы определить длины последующих последовательностей. Поскольку черновая модель запускается только на ведущем узле, её пакет настраивается последовательно и её выполнение начинается для заполнения её записей KV-кэша дополнительным токеном, созданным целью. Логиты, созданные этим черновым запуском, в текущем запуске будут использоваться для выборки первого чернового токена в следующем запуске. Самое главное, пока черновик работает, планирование следующего пакета может быть запущено.

Расписание MTP для одного токена

Хотя время выполнения ещё не предоставляет древообразное исследование черновика в стиле Eagle, мы реализовали специальный случай этой схемы, рассматривая линейную последовательность черновых токенов, сгенерированных моделью размером одного декодерного слоя трансформера. Эта схема может использоваться для прогнозирования черновиков с открытыми весами DeepSeek R1. Подслучай прогнозирования одного токена интересен, так как крупные слои MTP достигают достаточно высоких уровней принятия, чтобы оправдать их накладные расходы.

Планирование MTP несколько более сложное, так как черновая модель гораздо быстрее, скрывая меньшее количество задержек на стороне CPU. Кроме того, черновик разделён наряду с целевой моделью, требуя передачи общей памяти для информации о пакете. Запуск начинается с передачи информации о пакете и выборки первого токена из перенесенных логитов, подобно предыдущей схеме. Далее, целевая модель запускается для проверки токенов, обрабатывая 2 * D токенов, где D — это размер пакета декодирования. Это идеально для микропакетирования в моделях смеси экспертов (MoE) с медленными межсоединениями, такими как InfiniBand, так как пакет делится поровну на две половины. Скрытые состояния целевой модели переносятся в следующий черновик, в то время как логиты передаются в выборщик для верификации.

Выполняя ограниченное количество дополнительной работы на GPU, мы избегаем синхронизации CPU-to-GPU после принятия черновой последовательности. После того, как входные токены целей сдвинуты, ядро подключает следующие целевые токены в их соответствующие места. Черновик затем повторно запускается с той же информацией о пакете, что и цель, заполняя записи в KV-кэше и строя логиты и скрытые состояния для следующего запуска, выполняя некоторую избыточную работу на токенах, которые не были приняты. В этих ситуациях задержка неиспользуемой работы едва измеряется из-за малого размера черновой модели. Параллельно с запуском черновика, на CPU определяются длины последовательностей и начинается планирование следующего пакета, не нужно ждать завершения работы на GPU.

Накладные расходы дополнительной работы в слое черновика не заметны при внимании, однако слои MLP более проблематичны. Поскольку инструкции умножения матриц заполняются до границы в 64 вдоль измерения количества токенов, если удвоение не требует значительно больше блоков, накладные расходы скрыты. Для более длинных черновых последовательностей накладные расходы более значительны и схема, используемая для обычных моделей черновик-цель, работает лучше.

Ссылки

• Быстрый вывод из трансформеров через спекулятивное декодирование

• EAGLE: Спекулятивная выборка требует пересмотра неопределенности признаков

• EAGLE-2: Более быстрый вывод языковых моделей с экономными черновыми деревьями

• EAGLE-3: Увеличение ускорения вывода больших языковых моделей за счет тестов на этапе обучения

• Medusa: Простая платформа для ускорения вывода LLM с несколькими головами декодирования

• Нормализация слоя

• FlashInfer: Эффективный и настраиваемый механизм внимания для сервиса вывода LLM

• Технический отчет DeepSeek-V3