Оптимизация prompt lookup decoding в llama.cpp: ускорение до 140 раз
Hacker News · оригинал
Материал подготовлен автоматизированной редакционной системой. Факты можно сверить по указанному первоисточнику.
Разработчик представил набор оптимизаций для механизма предсказания токенов в llama.cpp, который ускоряет процесс черновика до 42 раз и снижает потребление памяти. Дополнительный вклад Дэниэля Лемира позволил достичь общего прироста скорости до 140 раз.
Улучшение эффективности локального инференса LLM критически важно для снижения задержек и требований к железу, что делает генерацию текста быстрее и доступнее на потребительских устройствах.
Механизм prompt lookup decoding (или n-gram speculation) используется в популярных движках инференса, таких как llama.cpp и vLLM, для ускорения генерации токенов. Этот метод является частным случаем спекулятивного декодирования, где роль черновой модели играет простая n-граммная модель. Алгоритм анализирует последовательность предыдущих токенов и предсказывает следующие на основе их частоты появления в корпусе данных. В llama.cpp для этого используются три типа кэшей: контекстный, динамический и статический, каждый из которых хранит статистику следования токенов за определенными n-граммами.
Исходная реализация в llama.cpp использовала вложенные структуры std::unordered_map, что приводило к неэффективному использованию памяти и времени. Автор статьи выявил, что внутренние карты копировались unnecessarily на каждом шаге черновика. Замена копирования на чтение по ссылке дала немедленный прирост скорости от 4,5 до 25,6 раз в зависимости от размера корпуса. Это базовое исправление стало отправной точкой для дальнейших оптимизаций.
Далее внешняя карта была заменена на ankerl::unordered_dense, а именно на вариант segmented_map. Это решение позволило ускорить загрузку статического кэша на 1,41–1,65 раза и снизить потребление памяти. Использование segmented_map вместо стандартного map-варианта помогло избежать пикового роста памяти при удвоении векторов, что особенно важно для больших корпусов данных.
Анализ распределения n-грамм показал, что большинство из них имеют лишь одного последующего токена. Это сделало использование хеш-карт для внутренних структур избыточным. Вместо этого автор предложил использовать отсортированные std::vector, что позволило сохранить логарифмическую сложность поиска через бинарный поиск. Специальная реализация цикла бинарного поиска, отделяющая изменение длины от сравнения значений, позволила CPU предугадывать следующие итерации, что ускорило черновик на 2,09 раза без статического кэша.
Ключевым шагом стало применение constmap от Дэниэля Лемира, построенного на binary fuse filters. Поскольку статический кэш неизменяем после загрузки, он идеально подходит для такой структуры. Данные упаковываются в непрерывный массив пар (токен, счетчик), а constmap хранит указатель на начало массива и количество элементов в одном 64-битном значении. Это позволило ускорить загрузку кэша в 6,32–16,12 раза и снизить пиковое потребление памяти до уровня, близкого к размеру самого файла.
Совокупность этих оптимизаций привела к тому, что черновик токенов стал до 42 раз быстрее, а потребление памяти снизилось в 2,6 раза. После того как Дэниэль Лемир внес свой PR, общий прирост скорости достиг 140 раз. Эти изменения не влияют на алгоритмическую логику принятия решений, но значительно улучшают производительность и эффективность использования ресурсов.