Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Неиспользованные варианты следующих токенов можно проверить вместе и получить больше готового текста за один проход целевой модели. В препринте Jungseob Lee и Sugyeong Eo, который не прошёл рецензирование, а числа в нём получили сами авторы, CAST ускорил декодирование относительно стандартной цепочки DFlash до 43%. Команде не нужно переобучать модель, но ширину дерева придётся подбирать под конкретную GPU и нагрузку.
CAST возвращает в работу отброшенные варианты
Опережающее декодирование (speculative decoding) делит генерацию между черновой и целевой моделями. Черновая модель дёшево предлагает продолжение, а целевая проверяет несколько токенов параллельно. Если предложение совпало с выбором целевой модели, система пропускает несколько последовательных шагов генерации.
DFlash составляет блок из 16 будущих позиций за один проход. Для каждой позиции он уже вычисляет оценки разных токенов, но стандартный декодер оставляет только лучший вариант и собирает из таких вариантов одну цепочку. Если целевая модель отклоняет токен в начале цепочки, остальные вычисленные варианты пропадают.
CAST сохраняет альтернативы и строит из них дерево. Узел соответствует возможному продолжению, а его оценка равна произведению оценок токенов на пути от корня. Алгоритм добавляет узлы в порядке убывания этой оценки и всегда сохраняет родительские узлы, поэтому любой выбранный путь остаётся полноценной последовательностью.
Целевая модель проверяет всё дерево одним проходом. Маска внимания разрешает каждому токену видеть общий исходный текст и только собственных предков, а не соседние ветви. Дополнительных запусков DFlash не возникает: цена расширения дерева складывается из упаковки данных, работы с кешем и более дорогой проверки целевой моделью.
CAST не меняет веса моделей и правило выбора токенов. Для жадной генерации дерево воспроизводит продолжение целевой модели, а для генерации со случайным выбором сохраняет её распределение ответов. При вычислениях с пониженной точностью форма пакета всё же может изменить выбор между почти равными токенами; такое же расхождение возникает у стандартного DFlash.
Широкое дерево помогает, пока проверка остаётся дешёвой
Максимальное дерево не всегда даёт максимальную скорость. Каждый новый узел повышает шанс принять более длинное продолжение, но одновременно увеличивает время прохода целевой модели. На некоторых GPU стоимость растёт плавно, а на других резко меняется после перехода к другой конфигурации вычислительного ядра.
CAST измеряет, сколько времени целевая модель тратит на дополнительный упакованный токен. Затем алгоритм добавляет следующий узел, пока его оценка вероятности принятия превышает объём текста, который система успела бы сгенерировать за добавленное время. В правой части этого сравнения находятся текущая скорость в токенах за миллисекунду и измеренная стоимость одного узла.
Такой критерий заменяет полный перебор ширины коротким замером задержки. Если проверка на сервере почти не дорожает, дерево расширяется. Если стоимость резко выросла, расширение прекращается, даже когда у черновой модели остались правдоподобные варианты.
CAST обошёл стандартную цепочку во всех восьми сочетаниях моделей и оборудования, а средний выигрыш составлял 20–36%. Порог стоимости оказался существеннее самой идеи широкого дерева: в конфигурации со скачком задержки дерево на 128 токенов ускоряло вывод лишь на 2%, тогда как выбранная CAST ширина давала 20%.
Под серверной нагрузкой тот же принцип распределяет общий бюджет проверки между запросами. Узлы разных деревьев попадают в одну очередь по оценке, поэтому запрос с более вероятным продолжением может получить больше места. Фиксированная ширина здесь проигрывает: с ростом параллелизма стоимость проверки меняется, и дерево, выгодное для одиночного запроса, становится лишней работой.
В план сервера добавляется профиль проверки, а не переобучение
Работа меняет планы команд, которые самостоятельно управляют контуром вывода модели и уже используют блочную черновую модель. Вместо выбора ширины по чужому бенчмарку нужно измерить задержку целевой модели на своём оборудовании, с обычной длиной контекста и рабочим параллелизмом. Этот профиль становится частью конфигурации декодера наряду с размером пакета и политикой очереди.
Интеграция затрагивает упаковку дерева, маску внимания, обход принятых ветвей и обновление кеша. Переобучать DFlash или целевую модель не требуется. Для закрытого API, где нельзя управлять проходами модели и кешем, CAST неприменим на стороне клиента.
Правило нельзя напрямую переносить на черновые модели, которые строят дерево последовательными проходами. Там каждый новый уровень увеличивает не только проверку, но и стоимость подготовки вариантов. CAST рассчитан на случай, когда один проход уже выдал оценки всех позиций блока.
Основные замеры охватывают две семьи моделей, три поколения GPU и пять областей с задачами по математике, коду и диалогам. DFlash проверяли непосредственно по времени декодирования, а перенос на второй однопроходный черновик DART — главным образом через повторное воспроизведение сохранённых оценок. Поэтому практический вывод относится прежде всего к архитектуре DFlash и близким к ней черновикам.
CAST не делает широкие деревья универсально выгодными. Он показывает более полезное правило проектирования: ширина — параметр развёртывания, а не свойство модели. Если стоимость проверки зависит от GPU и числа одновременных запросов, фиксировать её на этапе разработки нельзя.
Источники
Иллюстрация: рисунок из статьи «CAST: Cost-Aware Speculative Trees from One-Pass Block Drafters», Jungseob Lee, Sugyeong Eo, CC BY 4.0
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
