Назад к блогу

Как устроена предварительная выборка промптов в llama.cpp: кэш n-грамм и его оптимизация

Как устроена предварительная выборка промптов в llama.cpp: кэш n-грамм и его оптимизация

В llama.cpp спекулятивное декодирование может опираться не на отдельную черновую модель, а на статистику n-грамм текущего промпта — механизм prompt lookup. Статья разбирает, как из счётчиков коротких последовательностей токенов рождается черновик, какие пороги отсекают ненадёжных кандидатов и зачем создатели ушли от хеш-таблиц к бинарному формату.

Спекулятивное декодирование ускоряет генерацию за счёт того, что дешёвый источник предлагает несколько токенов вперёд, а основная модель их проверяет одним проходом. В llama.cpp один из таких источников — механизм prompt lookup. Он не требует отдельной черновой модели и опирается только на статистику коротких последовательностей токенов. Разобраться стоит в том, как именно из этой статистики получается черновик, почему структура кэша выглядит именно так и что даёт переход от хеш-таблиц к бинарному формату.

Что такое prompt lookup и какие типы спекуляции есть

n-грамма — это последовательность из n токенов. Кэш n-грамм ведёт статистику коротких n-грамм, а черновик вычисляется по вероятностям, выведенным из этой статистики. Для текущего промпта хранятся n-граммы размеров от 1 до 4; кэш обновляется по мере генерации новых токенов.

Тип спекуляции задаётся через --spec-type. Варианты draft-* опираются на отдельную черновую модель, а ngram-* строят черновик по статистике n-грамм текущего промпта.

N-граммные варианты различаются принципом поиска. ngram-simple ищет предыдущий совпадающий n-грамм и вставляет следующий m-грамм. ngram-map-k делает то же, но использует внутреннюю хеш-таблицу n-граммов в текущем контекстном окне. ngram-map-k4v — вариант с n-граммными ключами и до четырёх m-граммных значений, помечен как экспериментальный. ngram-mod использует хеш-пул, общий для всех серверных слотов, где отображение идёт от хеша n-грамма к следующему токену, а не к следующему m-грамму. ngram-cache выполняет поиск по n-граммному кешу через common_ngram_cache_draft с контекстным, динамическим и статическим кешами.

Устройство кэша n-грамм

Внешняя карта common_ngram_cache сопоставляет каждую n-грамму внутренней карте common_ngram_cache_part, которая хранит счётчики каждого токена словаря, следующего за данной n-граммой. Например, для n-грамма ("of", "the") внутренняя часть содержит { "city": 6, "war": 3, "year": 1 } — сколько раз каждый токен-кандидат встречался после этого контекста. По ключу-контексту внешняя карта находит внутреннюю, а по токену-кандидату во внутренней извлекается его счётчик.

Тип определён так:

typedef std::unordered_map<common_ngram, common_ngram_cache_part, common_ngram_hash_function> common_ngram_cache;

Изначально обе карты были std::unordered_map, но затем внешняя была заменена на ankerl::unordered_dense::segmented_map, а внутренняя — на отсортированный вектор пар (токен, счётчик) с отдельной картой смещений.

Кэш контекста хранит n-граммы размеров от 1 до 4 для текущих обрабатываемых токенов, динамический кэш — счётчики n-грамм из предыдущих запусков, статический — n-граммы размера 2 из статического корпуса. Для каждого токена-последователя во внутренней карте накапливается число его появлений после соответствующей n-граммы.

Как выбирается токен-кандидат и что задают пороги

Для каждого n llama.cpp берёт токен с наивысшим счётом $y^* = \arg\max_y s_n^{f}(y)$. Пусть $F(X_n) = \sum_y f(X_n, y)$ — число раз, когда $X_n$ встречалась с каким-либо токеном после неё. Кандидат принимается при выполнении двух условий:

$$F(X_n) \ge a_n \,\, \text{and} \,\, f(X_n, y^*) \ge p_n \, F(X_n)$$

То есть n-грамма $X_n$ должна появиться не менее $a_n$ раз, и токен $y^*$ должен следовать за ней в не менее чем доле $p_n$ от этих появлений.

Пороги различаются по кэшам. Для контекстного: $(a_1, a_2, a_3, a_4) = (2, 2, 1, 1)$ и $(p_1, p_2, p_3, p_4) = (0.66, 0.5, 0.5, 0.5)$. Для динамического: $(a_1, a_2, a_3, a_4) = (4, 3, 2, 2)$ и $(p_1, p_2, p_3, p_4) = (0.75, 0.66, 0.66, 0.66)$. Сначала оценка идёт по контекстному кэшу; к динамическому llama.cpp обращается, только если ни один кандидат из контекстного не прошёл ни для какого n. Если и там нет проходящего кандидата, происходит откат к статическому кэшу — уже как к единственному источнику, а не для перевзвешивания кандидатов. Если и статический не срабатывает, следующий токен не черновится.

Статический кэш и его сериализация

Статический кэш хранит отсортированный массив pair, где каждый элемент — пара из токена и числа, например ("city", 6), ("war", 3), ("year", 1). В constmap хранятся ключи из двух токенов, например ("of", "the"), и соответствующее им 64-битное значение: старшие 40 бит — позиция в массиве pairs, младшие 24 бита — количество пар.

Файл статического кэша хранит небольшой заголовок, массив pairs и сериализованный constmap последовательно. При загрузке файл читается в один буфер, а constmap открывается внутри буфера через fcm_verified_constmap_view.

Поиск common_ngram_cache_static_find берёт ключ из токенов n-грамма и выполняет один поиск в constmap через fcm_verified_constmap_lookup с размером ключа STATIC_KEY_SIZE. Если ключ не найден, кандидат отсутствует. Иначе из значения извлекаются позиция (сдвигом вправо на STATIC_LEN_BITS) и количество (маской STATIC_LEN_MASK), после чего берутся записи со смещением position и count. Поскольку pairs имеют тот же layout, что и отсортированные векторы, для поиска количества токена-кандидата используется бинарный поиск фиксированной длины.

Почему выбран бинарный формат, а не хеш-таблица

Причина в распределении последователей. 64% 2-грамм, используемых для черновиков статического кэша, имеют только одного последователя, поэтому std::vector для внутренней карты существенно экономичнее хеш-таблиц. Но распределение имеет тяжёлый хвост: у редких n-грамм тысячи последователей, и обычный вектор резко увеличил бы задержку поиска. Поэтому используется отсортированный вектор — поиск остаётся $\mathcal{O}(\log n)$ для n-грамм на хвосте.

Эффект на память: статический кэш теперь занимает примерно столько же памяти, сколько его файл — 463 MB для файла 467 MB. Пиковое потребление памяти снижается до 1.30x, с 1.71 GB до 1.31 GB на корпусе 541 MB. Загрузка статического кэша ускоряется в 6.32–16.12 раз, с 3.76 s до 0.23 s на том же корпусе. Черновики (drafting) со статическим кэшем ускоряются в 1.06–1.20 раз, а acceptance rate идентичен варианту с отсортированными векторами на каждом корпусе.

Бинарный поиск: фиксированное число итераций против зависимого

В одной реализации half вычисляется как n / 2, base обновляется в зависимости от сравнения base[half].first < token — либо base + half, либо base, — а n уменьшается на half. Здесь n уменьшается на одну и ту же величину независимо от результата сравнения, поэтому поиск по 8 элементам всегда занимает 3 итерации, с n от 8 до 4, до 2, до 1.

В другой реализации half = len / 2, mid = first + half, и при mid->first < token выполняется first = mid + 1 и len -= half + 1, иначе len = half. Здесь число итераций зависит от токена: поиск по 8 записям занимает 3 или 4 итерации.

Разница в том, как конвейеризуется поиск. В зависимой ветви CPU не может вычислить len != 0, пока запись текущей итерации не придёт из памяти. В ветви с фиксированным шагом от записи зависит только base, а n не зависит ни от какой записи, поэтому CPU может вычислить n > 1, не ожидая запись, и перейти к поиску следующего кандидата, пока чтения текущего ещё в полёте.

Как формируется черновик для ngram-cache

Во входной список копируется prompt, затем добавляется dparams.id_last, и в result кладётся dparams.id_last. После этого вызывается common_ngram_cache_draft с n_draft и границами LLAMA_NGRAM_MIN/LLAMA_NGRAM_MAX и тремя кэшами — контекстным, динамическим, статическим. Число запрашиваемых токенов задаётся переменной n_draft, которая инициализируется как 8, а не значением --spec-draft-n-max. После вызова, если result не пуст, из него удаляется первый токен (id_last).

accept. У ngram-cache этот метод noop. Общий common_speculative_accept лишь учитывает статистику и вызывает impl->accept(seq_id, n_accepted, false).

Как устроены другие n-граммные реализации

Для ngram-simple начало работы и учёт принятых токенов ничего не делают, обработка батча пропускается, а черновик для каждой последовательности строится по её промпту и последнему токену. Поведение настраивается тремя параметрами: size_n, size_m, min_hits.

Для ngram-map-k и ngram-map-k4v используется один класс: начало работы строит внутреннюю хеш-таблицу n-граммов текущего контекстного окна, обработка батча пропускается, черновик строится по промпту и последнему токену, а учёт принятых токенов пропускается для чужих последовательностей и обновляет состояние для своих. Поведение настраивается size_key, size_value, key_only и min_hits; режим key_only включается только для ngram-map-k.

Для ngram-mod начало работы сбрасывает позицию последнего токена и длину прошлого черновика, добавляет в общий хеш-пул все n-граммы промпта и при заполнении пула выше 0.25 очищает его. Черновик строится, только если промпт не короче длины n-грамма: новые n-граммы добавляются чанками, результат собирается из n-1 предыдущих токенов и последнего токена, после чего до n_max раз запрашивается следующий токен; если пул не даёт токен и текущая позиция меньше n_min, результат очищается, иначе обрезается по достигнутой длине. Учёт принятых токенов пропускается для чужих последовательностей, а для своих при ненулевой длине прошлого черновика считается доля принятых токенов, и при доле ниже 0.25 и пяти подряд низких приёмов пул очищается. Поведение настраивается n_match, n_max и n_min; при n_match меньше 16 выдаётся предупреждение о возможном низком качестве.

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

Счётчики означают следующее:

  • #calls(b,g,a) — число вызовов begin (новый промпт), generation и accumulation для данной реализации;
  • #gen drafts — число черновиков, сгенерированных реализацией;
  • #acc drafts — число черновиков, принятых (частично) основной моделью;
  • #gen tokens — число токенов, сгенерированных реализацией, включая отклонённые;
  • #acc tokens — число токенов, принятых основной моделью.

В коде они печатаются из полей n_call_begin, n_call_draft, n_call_accept, n_gen_drafts, n_acc_drafts, n_gen_tokens, n_acc_tokens. Draft acceptance rate выводится как отношение принятых к сгенерированным.

Пограничные случаи и компромиссы

При отсутствии совпадений llama.cpp последовательно спускается по уровням кэша и в худшем случае не формирует черновик вовсе. Для статического кэша токен берётся, если $C_{\text{st}}(X_2) \ge a_2 = 2$ и $c_{\text{st}}(X_2, y) \ge p_2 \, C_{\text{st}}(X_2) = 0.5 \, C_{\text{st}}(X_2)$.

Выбор n от 1 до 4 отражает компромисс: большие n дают более специфичные совпадения, но требуют больше памяти и реже встречаются, а меньшие n чаще срабатывают, но менее надёжны. В заметках указано, что малые n не рекомендуются, MoE, требуют длинных черновиков, а для плотных моделей предлагается уменьшать --spec-ngram-mod-n-min и --spec-ngram-mod-n-max. Пример запуска: llama-server ... --spec-type ngram-mod --spec-ngram-mod-n-match 24 --spec-ngram-mod-n-min 48 --spec-ngram-mod-n-max 64. Сказано лишь, что --spec-draft-n-max ограничивается обученным размером блока черновой модели.

Как сравнивали и что получилось

Стенд и методология. Все эксперименты выполнены на Apple M4 Pro с 14 ядрами и 48 ГБ памяти. Статические кэши строились утилитой llama-lookup-create на корпусе WikiText-103, затем тестовый текст WikiText-103 прогонялся через llama-lookup-stats. Помимо полного корпуса около 541 MB, кэши строились из первых 25, 50, 100 и 200 MB обучающего текста WikiText-103; размер корпуса 0 означает работу без статического кэша, что измеряет только контекстный и динамический кэши. Бенчмарк предполагает размер контекста модели 4096 токенов. Все результаты — медиана 3 прогонов, планки ошибок показывают минимум и максимум.

Результаты. По сценарию со статическим кэшем на корпусе 541 MB:

МетрикаЗначение
Загрузка статического кэша6.32x–16.12x быстрее, с 3.76 s до 0.23 s
Память статического кэша463 MB для файла 467 MB
Пиковая памятьдо 1.30x ниже, с 1.71 GB до 1.31 GB
Черновики (drafting)1.06x–1.20x быстрее
Acceptance rateидентичен варианту с отсортированными векторами на каждом корпусе

Оптимизация Daniel Lemire ускоряет drafting до 4.2x со статическим кэшем и до 1.9x без статического кэша поверх исходных оптимизаций.

По реализациям: для ngram_simple acceptance rate = 0.57576 (171 accepted / 297 generated); для ngram_mod acceptance rate = 0.70312 (90 accepted / 128 generated); для ngram_map_k — #calls(b,g,a) = 6 1690 26, #gen drafts = 26, #acc drafts = 26, #gen tokens = 1248, #acc tokens = 968.

Оговорки. Автор измерял три метрики, которые реально меняются от его правок: задержку на один сгенерированный токен, время загрузки статического кэша и память, занимаемую статическим кэшем. Сам алгоритм prompt lookup decoding он не менял, поэтому датасет влияет в основном на acceptance rate, который его изменения не затрагивают. Стенд ограничен одной конфигурацией — Apple M4 Pro с 14 ядрами и 48 ГБ памяти — и фиксированным размером контекста модели 4096 токенов. Каждое число — медиана трёх прогонов с усами min/max. End-to-end эффект спекулятивного декодирования (throughput, latency, draft acceptance) измеряется отдельно через SPEED-Bench; приведённые цифры его не показывают.

Что из этого следует на практике

  • Черновик — это не алгоритмическая догадка, а статистика. Токен-кандидат выбирается по счётчикам из кэша, и его принятие зависит от двух порогов: минимального числа вхождений n-граммы и минимальной доли, в которой кандидат за ней следовал. Меняя эти пороги, вы меняете соотношение между числом черновиков и их точностью.
  • Порядок кэшей определяет, что вы получите. Контекстный кэш оценивается первым, динамический — только при полном провале контекстного, статический — как последний источник. Поэтому статический кэш работает не как перевзвешивание, а как запасной вариант.
  • Формат статического кэша — компромисс под распределение данных. Большинство 2-грамм имеют одного последователя, поэтому вектор экономичнее хеш-таблиц, но тяжёлый хвост требует логарифмического поиска. Отсюда и бинарный формат: файл читается в один буфер, а constmap открывается прямо внутри него.
  • Число итераций бинарного поиска влияет на конвейер. Когда длина поиска зависит от результата сравнения, CPU ждёт запись из памяти; когда шаг фиксирован, он может продолжить поиск следующего кандидата, пока чтения ещё в полёте.
  • Размер черновика для ngram-cache задан в коде, а не флагом. n_draft инициализируется как 8, а не значением --spec-draft-n-max; для ngram-mod длина управляется через n_match, n_max, n_min, и при n_match < 16 выдаётся предупреждение о возможном низком качестве.

Где смотреть в коде

Источники

Похожее