Разреженное префиксное кеширование, или где расставлять чекпойнты в гибридных LLM

У нас с моим аспирантом Михаилом Широких на днях вышла статья Sparse Prefix Caching for Hybrid and Recurrent LLM Serving. В этом посте я постараюсь популярно рассказать, о чём она, да и вообще дать общее введение в тему.

Префиксное кеширование: общая идея

Когда современная LLM обрабатывает запрос, основная вычислительная работа делится на две фазы. Первая — префилл (prefill): модель один раз читает весь входной промпт и формирует внутреннее состояние. Вторая — декодирование (decoding): порождение ответа токен за токеном. Префилл может занимать значительную часть общего времени, особенно если контекст длинный, а в современных моделях, как все мы знаем, он легко съедает сотни тысяч токенов, и даже миллион может переполниться.

Теперь представим, что приходит несколько запросов с одним и тем же длинным началом. Например:

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

В таких ситуациях глупо каждый раз заново прогонять одинаковый огромный префикс через модель. Хочется один раз посчитать внутреннее состояние после длинного префикса и потом переиспользовать его — это и называется префиксное кеширование (prefix caching).

В классическом трансформере это работает достаточно естественным образом: внутреннее состояние модели на длинном префиксе — это, по сути, набор векторов ключей и значений (KV-кеш) для каждого токена и каждой головы внимания. Современные системы для использования LLM (LLM serving), в частности vLLM с PagedAttention, SGLang с RadixAttention, Preble и другие, построены вокруг того, чтобы хранить и переиспользовать эти KV-кеши максимально гибко: их разбивают на блоки, организуют в виде префиксных деревьев, передают между машинами и так далее.

А если у нас не трансформер?

Но в последние пару лет, как я не раз уже в своём блоге обсуждал, всё более важную роль играют архитектуры, которые либо полностью отказываются от внимания (Mamba, RWKV и другие state-space модели), либо комбинируют его с рекуррентными или SSM-слоями — это так называемые гибридные модели. Например, у Qwen-3.5, которым мы пользовались в экспериментах в этой статье, часть слоёв — обычное внимание, а часть — GatedDeltaNet, рекуррентный слой со скрытым состоянием.

Но рекуррентный слой устроен принципиально иначе! Вместо того, чтобы хранить отдельные пары (ключ, значение) для каждого токена, он поддерживает скрытое состояние фиксированного размера, которое обновляется на каждом шаге:

$h_t = F(h_{t-1}, x_t).$

Таким образом, чтобы продолжить вычисление с позиции $t$ , нужно знать только $h_t$ , а не всю историю состояний $h_1, \ldots, h_t$ . Это значит, что для рекуррентных слоёв “честный” KV-кеш в стиле трансформера — это, во-первых, избыточно (нам не нужны все промежуточные состояния), а во-вторых, в случае некоторых архитектур ещё и очень дорого: например, у GatedDeltaNet состояние на одну “голову” имеет размер $d_{\text{head}}^2$ , тогда как в обычном внимании оно линейно по $d_{\text{head}}$ .

Получается, что для гибридных моделей выбор “либо хранить плотный KV-кеш на каждый токен, либо вообще ничего не хранить и прогонять префикс заново” — какой-то заведомо ложный и неоптимальный. У нас ведь есть и третий вариант: хранить точные рекуррентные состояния в нескольких опорных точках вдоль префикса. Назовём такие точки чекпойнтами (checkpoints).

Идея разреженного префиксного кеширования

Дальше собственно начинается суть работы Михаила. Идея простая: пусть мы закешировали префикс длины $N$ , и у нас есть бюджет из $M$ чекпойнтов (где $M \ll N$ ). Мы выбираем $M$ позиций $c_1 < c_2 < \ldots < c_M$ и сохраняем рекуррентное состояние именно в этих позициях.

Приходит новый запрос. Он совпадает с закешированным префиксом до какой-то глубины $t$ , а дальше расходится. Мы находим самый глубокий чекпойнт $c_i \leq t$ и:

восстанавливаем оттуда точное рекуррентное состояние $h_{c_i}$ ;
доигрываем рекуррентный слой по токенам $x_{c_i+1}, \ldots, x_t$ , то есть тратим $t - c_i$ шагов пересчёта.

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

Главный вопрос теперь становится алгоритмическим: где именно расставить $M$ чекпойнтов, чтобы минимизировать ожидаемые вычислительные затраты?

Где ставить чекпойнты?

Ответ зависит от того, что мы знаем о будущих запросах — точнее, о распределении глубин совпадения с закешированным префиксом. Обозначим через $p_t = \Pr[T = t]$ вероятность того, что следующий запрос совпадёт с префиксом ровно до позиции $t$ . Тогда ожидаемые затраты при наборе чекпойнтов $\mathcal{C}$ составят

$\mathbb{E}[r(T; \mathcal{C})] = \sum_{t=1}^{N} p_t \cdot (t - \ell(t; \mathcal{C})),$

где $\ell(t; \mathcal{C})$ — позиция ближайшего сверху чекпойнта (или ноль, если ни один не подходит). Получается одномерная односторонняя задача о взвешенных $k$ -медианах: расставить $M$ точек на отрезке так, чтобы минимизировать суммарное взвешенное расстояние от точек распределения до ближайшего чекпойнта слева.

На этом рисунке сверху распределение $p(t)$ глубин совпадения будущих запросов с закешированным префиксом (синтетическое, для иллюстрации), а снизу — функция $r(t) = t - \ell(t; \mathcal{C})$ , которая показывает, сколько токенов придётся пересчитать, если совпадение случилось до глубины $t$ ; Маркеры ▼ — позиции чекпойнтов при равномерной расстановке:

На всех таких картинках есть характерная “пила” $r(t)$ : после каждого чекпойнта функция обнуляется и линейно растёт до следующего, ведь если совпадение оборвалось через $k$ токенов после ближайшего чекпойнта слева, надо пересчитать ровно эти $k$ токенов. А ожидаемая стоимость $\sum_t p_t \cdot r(t)$ — это просто интеграл произведения двух кривых, верхней и нижней.

Если зуб пилы попадает в место, где $p(t) \approx 0$ (как, например, в долине $t \in (75, 110)$ на картинке), то этот зуб может быть очень высоким, но он ничего не стоит — туда никакие запросы всё равно не приходят. А вот высокий зуб, попавший в пик распределения, дорого обойдётся. Отсюда и оптимизационная задача: расставить чекпойнты так, чтобы зубы были низкими там, где у $p(t)$ сосредоточена масса.

Случай 1: когда про распределение мы ничего не знаем

Если распределение $p_t$ равномерное, или если мы хотим оптимизировать в худшем случае, то ответ оказывается очень простым и интуитивным: чекпойнты нужно расставлять равномерно. Сбалансированное расположение, где все промежутки между соседними чекпойнтами отличаются друг от друга не больше чем на единицу, оптимально и в среднем под равномерным распределением, и в худшем случае.

Доказательство здесь устроено довольно симпатично. Обозначим длины промежутков между чекпойнтами через $g_i = c_{i+1} - c_i$ (где $c_0 = 0$ , $c_{M+1} = N+1$ ); это положительные целые числа, в сумме дающие $N+1$ . На промежутке длины $g_i$ затраты на пересчёт принимают в точности значения $0, 1, \ldots, g_i - 1$ (по одному на каждую позицию внутри промежутка). При равномерном распределении $p_t = 1/N$ каждая позиция вносит одинаковый вклад, и ожидаемый пересчёт получается равным

$\mathbb{E}[r(T; \mathcal{C})] = \frac{1}{N}\sum_{i=0}^{M} \frac{g_i(g_i - 1)}{2} = \frac{1}{N}\sum_{i=0}^{M} \phi(g_i),$

где $\phi(x) = x(x-1)/2$ — строго выпуклая функция. Дальше работает классический аргумент про выпуклость: если найдутся два промежутка с $g_i \geq g_j + 2$ , то замена $(g_i, g_j) \to (g_i - 1, g_j + 1)$ сохраняет сумму и строго уменьшает целевую функцию (ведь разность равна $g_i - (g_j - 1)> 0$ ). Значит, в оптимуме все промежутки отличаются не больше чем на единицу.

В численном виде это выглядит так: положим $K = M+1$ и запишем $N+1 = qK + \rho$ , где $0 \leq \rho < K$ . Тогда оптимум ровно один (с точностью до перестановки): $K - \rho$ промежутков длины $q$ и $\rho$ промежутков длины $q+1$ . Подставляя обратно, получаем

$\mathbb{E}[r(T; \mathcal{C}^\star)] = \frac{N}{2(M+1)} + O(1).$

Для худшего случая работает ровно тот же аргумент : $\max_t r(t; \mathcal{C}) = \max_i g_i - 1 \geq \lceil (N+1)/(M+1)\rceil - 1$ по принципу Дирихле, и сбалансированная расстановка достигает нижней границы.

Важный частный случай здесь возникает, когда $M = O(\sqrt{N})$ : и число чекпойнтов, и средние/худшие затраты на пересчёт растут как $\sqrt{N}$ , то есть память и вычисления находятся в естественном равновесии.

Но в этом, конечно, ничего нового нет, интересные вещи начинаются дальше.

Случай 2: распределение неравномерное

В реальности распределения далеко не равномерные. Например, у документов из датасета QuALITY есть резкий пик в районе полной длины документа, потому что вопросы обычно задаются после всего текста. У System Prompts (датасет, собственно, системных промптов) масса сосредоточена ближе к короткому общему началу, потому что пользовательские запросы быстро расходятся; а у NarrativeQA, наоборот, мультимодальное широкое распределение.

В таком случае равномерное размещение неоптимально: лучше “сгущать” чекпойнты там, где у распределения большая масса. Для этой версии задачи мы выписываем точное решение динамическим программированием. Но я всё-таки распишу это подробно, потому что здесь есть пара важных особенностей.

Пусть $\mathrm{dp}[m, j]$ — это минимальная ожидаемая стоимость обслуживания глубин ${1, \ldots, j}$ при бюджете $m$ чекпойнтов. Если последний чекпойнт стоит в позиции $s \leq j$ , то стоимость распадается на две части: оптимальное обслуживание глубин ${1, \ldots, s-1}$ с $m-1$ чекпойнтами (это $\mathrm{dp}[m-1, s-1]$ ) и стоимость пересчёта от чекпойнта $s$ до конца, $w(s, j) = \sum_{t=s}^{j} p_t (t - s)$ . Минимизируем по выбору $s$ и получаем стандартную рекуррентность:

$\mathrm{dp}[m, j] = \min_{1 \leq s \leq j} \big(\mathrm{dp}[m-1, s-1] + w(s, j)\big).$

В лоб это заняло бы $O(N^2 M)$ : каждый из $NM$ элементов таблицы требует перебора $O(N)$ кандидатов $s$ . Но дальше можно посмотреть на структуру задачи внимательнее.

Введём префиксные суммы $P_j = \sum_{t \leq j} p_t$ и $T_j = \sum_{t \leq j} t \cdot p_t$ . Тогда $w(s, j) = (T_j - T_{s-1}) - s(P_j - P_{s-1})$ , и наша рекуррентность переписывается в виде

$\mathrm{dp}[m, j] = T_j + \min_{s \leq j} \big[(-s) \cdot P_j + b_s\big], \qquad b_s = \mathrm{dp}[m-1, s-1] - T_{s-1} + s P_{s-1}.$

То есть для фиксированного $m$ каждый кандидат $s$ соответствует прямой $y = (-s) \cdot x + b_s$ , а вычисление $\mathrm{dp}[m, j]$ сводится к нахождению её нижней огибающей в точке $x = P_j$ . Здесь работает классический “трюк с выпуклой оболочкой” (convex hull trick): коэффициенты –s монотонно убывают по s, а точки $P_j$ монотонно растут по $j$ , ведь $p_t \geq 0$ . Это значит, что нижнюю огибающую можно строить инкрементально, а указатель ответа в ней движется только вперёд: каждый слой DP считается за амортизированное время $O(N)$ , и всё решение требует только $O(NM)$ .

Вычисление расстановки достаточно делать оффлайн и пересчитывать только при существенном обновлении эмпирической гистограммы, так что в рамках собственно inference такой алгоритм не добавляет никакого оверхеда.

Вот неравномерный пример (смесь двух гауссианов) и две стратегии размещения шести чекпойнтов, равномерная (оранжевые маркеры) и оптимальная (синие):

Как видите, оптимальная расстановка кучкуется вокруг двух пиков распределения, оставляя долину между ними, где $p(t) \approx 0$ , без покрытия, и ожидаемые затраты падают с 14.2 до 8.6 (на 39%) при том же бюджете чекпойнтов.

Вот для полноты картины развёрнутый пример с тремя распределениями; тут видно, как DP ставит чекпойнты именно там, где сосредоточена масса:

Этот рисунок показывает и логичную закономерность: чем менее равномерным будет распределение и чем меньше бюджет, тем сильнее выигрыш DP над равномерной расстановкой. Это интуитивно понятно: при большом бюджете и равномерная стратегия в какой-то момент покроет все важные регионы, а при сильно неравномерном распределении любой впустую потраченный чекпойнт особенно дорог. Эта закономерность будет хорошо видна и в основных экспериментах ниже.

Случай 3: оцениваем распределение по истории

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

Это даёт нам доказанная Михаилом лемма об устойчивости: значение функции $\mathbb{E}_p[r(T; \mathcal{C})]$ липшицево по $p$ в норме полной вариации, то есть

$\big|\mathbb{E}_p[r(T; \mathcal{C})] - \mathbb{E}_q[r(T; \mathcal{C})]\big| \leq N \cdot |p - q|_1.$

Доказательство короткое и изящное: разность ожиданий — это $\sum_t (p_t - q_t) \cdot r(t; \mathcal{C})$ , и поскольку $0 \leq r(t; \mathcal{C}) \leq N$ , по неравенству треугольника получаем нужную оценку. Лемма выглядит почти тривиально, но именно она позволяет нам получить всё остальное: из неё следуют два важных утверждения.

Подстановочная гарантия (plug-in guarantee). Пусть $\hat{p}^{(n)}$ — эмпирическая гистограмма после $n$ независимых наблюдений из истинного распределения $p$ , и $\widehat{\mathcal{C}}_n$ — оптимальная расстановка, посчитанная по $\hat{p}^{(n)}$ нашим динамическим программированием. Тогда с вероятностью не меньше $1 - \delta$ выполняется

$|\hat{p}^{(n)} - p|_1 \leq \sqrt{N/n} + \sqrt{2\log(1/\delta)/n}.$

Доказательство тоже короткое: математическое ожидание полной вариации ограничивается через неравенства Коши и Йенсена ( $\mathbb{E}|\hat{p}^{(n)} - p|_1 \leq \sqrt{N/n}$ ), а концентрация вокруг ожидания получается из неравенства Макдиармида, поскольку добавление одного наблюдения меняет полную вариацию не больше чем на $2/n$ .

Подставляя эту оценку в лемму об устойчивости, получаем

$0 \leq \mathbb{E}_p[r(T; \widehat{\mathcal{C}}_n)] - V_M(p) \leq 2N \big(\sqrt{N/n} + \sqrt{2\log(1/\delta)/n}\big),$

где $V_M(p)$ — оптимальная стоимость при истинном распределении. То есть отклонение от оптимума убывает как $1/\sqrt{n}$ .

Следование за дрейфом распределения. В реальных системах распределение глубин совпадения меняется со временем: меняются сценарии использования, растёт типичная длина контекстов, кеш насыщается и т.п. Для борьбы с этим мы используем классический приём — экспоненциально взвешенную гистограмму с весами $\propto \gamma^{t-s}$ :

$\hat{p}^{(\gamma)}_t = \frac{1-\gamma}{1-\gamma^t} \sum{s=1}^{t} \gamma^{t-s} \mathbf{e}_{T_s}.$

Здесь $\mathbf{e}_{T_s}$ — индикатор того, что наблюдалось ровно совпадение глубины $T_s$ . Для такой оценки получается естественное разложение типа bias-variance:

$\mathbb{E}\big[|\hat{p}^{(\gamma)}_t - p_t|_1\big] \leq {\frac{1-\gamma}{1-\gamma^t}\sum_{s=1}^{t} \gamma^{t-s} |p_s - p_t|_1} + {\sqrt{\frac{N(1-\gamma)(1+\gamma^t)}{(1+\gamma)(1-\gamma^t)}}}.$

При больших $t$ разброс выходит на асимптоту $\sqrt{N(1-\gamma)/(1+\gamma)}$ , а смещение при ограниченном дрейфе $|p_u - p_{u-1}|_1 \leq \Delta$ не превосходит $\Delta\gamma/(1-\gamma)$ . Выбор $\gamma$ балансирует эти два слагаемых: большие $\gamma$ дают меньший разброс, но хуже отслеживают дрейф. В наших экспериментах мы используем $\gamma = 0.99$ , и эмпирически результаты слабо чувствительны к точному значению $\gamma$ в окрестности этого выбора.

Получается, что теоретически всё вообще замечательно: оптимальное оффлайн-решение по точному распределению устойчиво и к ошибкам оценивания (через лемму об устойчивости + концентрацию эмпирической гистограммы), и к смещению распределения со временем (через bias-variance для экспоненциально взвешенной оценки).

Экспериментальные результаты

Чтобы это проверить, мы взяли три датасета с естественной структурой длинного общего префикса.

QuALITY — длинные документы, по каждому из которых задаётся серия вопросов.
NarrativeQA — то же самое, но с ещё более длинными документами (целые книги).
System Prompts — реальные системные промпты от больших провайдеров LLM, к которым приклеены реальные пользовательские запросы из ShareGPT.

Сравнивали несколько стратегий расстановки чекпойнтов: сбалансированное размещение, блочное (каждые $B$ токенов), $\sqrt{L}$ -размещение, логарифмическое (экспоненциально растущие промежутки) и наше DP-optimal по экспоненциально взвешенной эмпирической гистограмме.

Все методы использовались с фиксированной last-K политикой кеширования записей, чтобы изолировать именно эффект расстановки чекпойнтов внутри отдельной записи (здесь есть ещё отдельная смежная задача про политики admission/eviction в стиле метода Marconi, мы её не трогаем, а просто фиксируем).

Как выглядят реальные распределения

Прежде чем смотреть на расстановки, полезно посмотреть, как реально выглядит распределение глубин совпадения на наших датасетах:

Видно, что распределения совсем не равномерные и сильно отличаются друг от друга: у QuALITY резкий пик в районе полной длины документа (медиана 6683 токенов, и почти вся масса сосредоточена около этого пика), у NarrativeQA — широкое многомодальное распределение (медиана около 57k), а у System Prompts масса сидит ближе к началу префикса (медиана 3274). Это ровно тот случай, ради которого затевалась вся теория из случая 2: нет одной фиксированной формы, распределения неравномерные, и наивное равномерное размещение будет сильно проигрывать.

Экономия токенов на длинных последовательностях

NarrativeQA даёт самые длинные тексты — настолько, что наш прототип реализации на доступном железе с ними не справлялся. Зато для этого датасета можно симулировать пересчёт без реального запуска модели и измерить фактор сокращения (recurrent-work reduction factor), то есть во сколько раз пересчёт уменьшается по сравнению со стратегией без кеширования.

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

Реальные измерения времени (wall-clock)

На QuALITY и System Prompts последовательности короче, и можно измерить настоящее время прогона на железе. Wall-clock мы измеряли на репрезентативной группе слоёв (1 attention + 3 GatedDeltaNet) у Qwen-3.5-0.8B, что, конечно, не вполне полноценный production-бенчмарк, но уж что смогли:

На Парето-фронтире DP-optimal либо доминирует над всеми базовыми стратегиями, либо совпадает с лучшим baseline (как правило, это блочное кеширование) но имеет заметно меньшее числом чекпойнтов. На System Prompts, где распределение особенно концентрированное у начала, DP даёт самые большие выигрыши при малых бюджетах — ровно как и предсказывает наш анализ. На QuALITY разрыв меньше, потому что распределение почти точечное около полной длины документа: там и расставлять-то особо нечего, все стратегии быстро упираются в один и тот же пик.

Дополнительно мы проверили, что время инференса действительно почти линейно зависит от числа реально пересчитываемых токенов, с небольшой константой на перенос состояния с CPU на GPU.

Это означает, что фактор сокращения, который мы измеряли на NarrativeQA, действительно превращается в реальную экономию времени.

Как выглядят расстановки на реальных данных

Наконец, любопытно просто посмотреть, как именно DP расставляет чекпойнты на реальных диалогах.

На QuALITY DP-optimal располагает чекпойнты неравномерно, концентрируя их ближе к местам, где у конкретных запросов происходило совпадение в предыдущих раундах:

На System Prompts расстановка ещё более показательная: чекпойнты группируются у короткого начала, где сидит вся масса, а длинный хвост остаётся без покрытия — потому что туда почти никто не доходит.

Выводы и что дальше

Главное ограничение такого подхода — он имеет смысл, когда запросы делят длинный, но не полностью совпадающий префикс внутри одной закешированной записи. То есть это сценарии вроде “много вопросов по одному документу”, “общий длинный системный промпт + разные пользовательские запросы”, ну или RAG.

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

Что ещё важно, в отличие от методов сжатия KV-кеша, при всём этом мы не меняем сами вычисления модели — выход абсолютно идентичен честному префиллу. Соответственно, наш метод хорошо комбинируется со всеми существующими подходами к сжатию KV-кеша для attention-слоёв.

Каковы следующие шаги? Их много! И все надо будет попробовать.

Полноценная интеграция в production-рантайм. Сейчас наши эксперименты — это прототип; для серьёзных бенчмарков нужно реализовать эффективный API для извлечения и восстановления состояния в Flash Linear Attention или похожих библиотеках. Низкоуровневые ядра (kernels) уже возвращают рекуррентные состояния на нужных позициях, дело за обёрткой.
Совмещение с политиками admission/eviction в стиле Marconi: они решают, какие закешированные префиксы держать в кеше, а мы — где внутри них ставить чекпойнты. Эти две задачи, по идее, должны естественно дополнять друг друга, но мы пока не понимаем как именно.
Обобщение теории на префиксные деревья с ветвлением, а не на одиночный префикс. Для типичных топологий (звезда: один документ, много вопросов; гребёнка: длинный prompt, разные хвосты) задача распадается на независимые однопрефиксные подзадачи, но в общем случае это уже более сложная комбинаторика.

В общем, интересная наука, и я категорически желаю Михаилу дальнейших успехов! Впрочем, надеюсь, и эту конкретную статью мы ещё немножко доработаем и подадим в приличное место.

Сергей Николенко

P.S. Прокомментировать и обсудить пост можно в канале «Sineкура»: присоединяйтесь!

Разреженное префиксное кеширование, или где расставлять чекпойнты в гибридных LLM

Префиксное кеширование: общая идея

А если у нас не трансформер?

Идея разреженного префиксного кеширования

Где ставить чекпойнты?

Случай 1: когда про распределение мы ничего не знаем

Случай 2: распределение неравномерное

Случай 3: оцениваем распределение по истории

Экспериментальные результаты

Как выглядят реальные распределения

Экономия токенов на длинных последовательностях

Реальные измерения времени (wall-clock)

Как выглядят расстановки на реальных данных

Выводы и что дальше

Comments

Leave a Reply Cancel reply

More posts

PEPPERED: An Existential Platformer

Охота на снарка: доказательство гипотезы о двойном покрытии циклами

Train Valley Origins

Curyeux