Ускорение Sonar через спекуляцию

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

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

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

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

Развернуть

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

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

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

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

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

Существуют различные методы реализации спекулятивного декодирования. В этой статье мы сосредоточимся на схемах, которые мы использовали для ускорения моделей 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 затем упаковываются в буфер, транслируемый всем параллельным рангам, выполняющим модель. Наконец, метаданные копируются в память графического процессора, и модель выполняется для получения логитов, из которых выбирается следующий токен.

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

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

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

График работы целевой-вспомогательной модели

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

Развернуть

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

График MTP для одного токена

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

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

Развернуть

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

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

Ссылки

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

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

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

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

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

Нормализация слоев

FlashInfer: Эффективный и настраиваемый движок внимания для обслуживания вывода LLM

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

Цитирование

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}