Как удешевить inference: KV Cache, batching и quantization

Чтобы сгенерировать token, decoder-only transformer берёт всю последовательность на данный момент, прогоняет её через каждый слой и считывает распределение вероятностей с последней позиции. Затем добавляет выбранный token и делает это снова. Это описание корректно, и именно это делает медленный запуск.

Оно также чрезвычайно расточительно, и причина — causal mask из главы 9. Векторы key и value позиции 7 вычисляются из входа позиции 7 и позиций перед ней. Когда приходит позиция 8, позиция 7 не может её видеть — именно это и означает causal, — поэтому key и value позиции 7 являются ровно теми же числами, что и раньше. Медленный запуск всё равно пересчитывает их на каждом шаге.

Значит, сохраните их. Это хранилище и есть key-value cache, самая важная оптимизация в serving языковых model:

out = model(prompt_ids, use_cache=True)          # prefill: the whole prompt
past = out.past_key_values                        
nxt = out.logits[:, -1].argmax(-1, keepdim=True)

for _ in range(n - 1):
    out = model(nxt, past_key_values=past, use_cache=True)   
    past = out.past_key_values                                
    nxt = out.logits[:, -1].argmax(-1, keepdim=True)

Посмотрите, что подаётся в model внутри цикла: nxt, один token. Не последовательность. Query нового token attends against каждый cached key, а cached keys всё равно никогда не должны были измениться. Это не approximation — проверка идентичности вывода выше и есть суть. Cache не обменивает качество на скорость; он удаляет лишнюю арифметику.

Чтобы ясно увидеть scaling, уберите transformer и замерьте один attention head с $d = 64$, один шаг generation, посчитанный обоими способами:

Правый столбец — причина. Пересчёт строит полную $n \times n$ attention matrix на каждом шаге — ту самую $O(n^2)$ из блока asymptotic notation в главе 9, оплачиваемую один раз на token. С cache вы вместо этого строите строку $1 \times n$: при 4 096 token — 67 MB scores против 16 KB.

Если считать multiply-accumulates вместо миллисекунд, машина исчезает из аргумента. Чтобы сгенерировать $T$ token с cold start:

На каждом шаге версия с cache линейна по context, а без cache — квадратична; если суммировать по generation, получается $O(T^2)$ против $O(T^3)$, и отношение растёт без предела. Девятикратная разница в начале была измерена всего на 48 token — меньше первой строки этой таблицы.

Cache также меняет то, что должно находиться в памяти. На laptop GPU с 8 GB при generation 256 token в fp16, если взять пик allocator и вычесть resident weights:

В 8,3 раза больше памяти, потраченной на то, чтобы произвести те же token медленнее. Это обещание из главы 5, пришедшее с неожиданной стороны: там reverse-mode autodiff должен был держать каждый intermediate живым для backward pass, и activations доминировали в памяти обучения. В inference нет backward pass и нечего удерживать ради него — поэтому память вместо этого доминируется cache, и это осознанный выбор, а не неизбежная стоимость.

Снова посмотрите на быстрый запуск: его первый token вёл себя не так, как остальные сорок семь.

prefill, 40 prompt tokens : 1.0224 s   ->  25.6 ms per token
decode,  47 steps         : 0.1665 s mean per step

Prompt стоил 25,6 ms на token, а каждый сгенерированный token — 166 ms. Тот же model, то же hardware, те же weights, шестикратная разница на token — и в направлении, которого большинство не ожидает. Prompt — дешёвая часть. Generation распадается на две фазы с действительно разной физикой:

Обе половины измеримы. Prefill, один pass по $P$ token:

Decode, один token against cache из $C$:

Прочитайте вторую таблицу дважды. Переход от 16 token context к 1 024 — в шестьдесят четыре раза больше history для attention — не изменил стоимость шага измеримо. Attention against cache — реальная работа, но её затмевает фиксированная стоимость протаскивания полумиллиарда weights через memory bus, чтобы получить один vector. Эта фиксированная стоимость — причина всего в следующем разделе.

Эти две фазы — источник двух чисел, которые сообщает любая serving system. Time to first token — по сути prefill, и он растёт вместе с prompt, поэтому длинный разговор медленно стартует. Tokens per second — это $1/\text{decode step}$, и он примерно постоянен, поэтому ответ затем течёт ровно. Chat, который медленно начинается, а потом плавно streamится, — не трюк rendering. Это эти две таблицы.

Cache обменивает арифметику на память, и память, которую он хочет, не мала. Для каждого token в context каждый layer хранит один key vector и один value vector на key-value head:

2 — это keys и values; всё остальное — architecture. Для model, измеряемого во всей этой главе, — 24 layers, 14 query heads, 2 key-value heads, head dimension 64 — в fp16 это $2 \times 24 \times 2 \times 64 \times 2 = 12{,}288$ bytes на token.

Формулы в этой области имеют привычку ошибаться на factor of two, поэтому проверьте её по allocator, а не верьте на слово:

KV cache tensors per layer: (1, 2, 295, 64) float16
measured: 3,624,960 bytes for 295 tokens = 12,288 bytes/token
formula : 2 * 24 * 2 * 64 * 2                = 12,288 bytes/token

Точно, и остаётся точно для каждой проверенной shape:

На последние три строки стоит посмотреть ещё раз. Тридцать два пользователя по 2 048 token у каждого, шестьдесят четыре по 1 024, сто двадцать восемь по 512 — cache во всех случаях равен 768 MB, потому что во всех трёх находятся 65 536 token. Cache зависит только от общего числа resident token, а не от того, как они распределены между пользователями. Этот факт — фундамент раздела про batching.

Глава 9 представила multi-query и grouped-query attention и отложила причину до этой главы. Причина — эта формула, а точнее $H_{kv}$ в ней.

Стандартный multi-head attention даёт каждому query head собственные key и value heads. У этого model 14 query heads; с полным multi-head attention его cache был бы $2 \times 24 \times 14 \times 64 \times 2 = 86{,}016$ bytes на token — 84 KB вместо 12 KB, ровно в семь раз больше, по отношению query heads к key-value heads.

Multi-query attention¹ доводит это до предела: все query heads делят один key-value head. Grouped-query attention² — компромисс, который победил: несколько key-value heads, каждый разделяется группой query heads, — потому что потеря качества у MQA была реальной, а у GQA нет. Ни один из них не покупает арифметику. Они существуют, чтобы делить эту формулу на целое число, и распространились по индустрии в тот момент, когда длинные contexts сделали cache binding constraint.

А он делает это быстро. Для model класса 7B с 32 layers и 8 key-value heads размерности 128 cache составляет 128 KB на token в fp16:

Собственные weights этого model — 13,0 GB в fp16, число из таблицы в конце главы. Поэтому при context в 128 000 token cache одного пользователя больше, чем model. Эту арифметику глава 16 превращает в деньги, и именно поэтому длинный разговор не просто медленный — он занимает фиксированный кусок машины всё время, пока жив запрос.

Decode memory-bound: weights протаскиваются через bus, чтобы произвести один token, а арифметические блоки простаивают. Значит, добавьте больше работы в тот же шаг. Запустите несколько запросов одновременно, и weights, прочитанные один раз, обслужат их всех. Измерено на том же model: каждый запрос держит cache на 64 token и decode один token:

Читайте два правых столбца вместе, потому что в них весь смысл. Переход от одного запроса к шестнадцати умножает throughput на 6,0 и умножает ожидание любого отдельного запроса на 2,67. Batch сделал сервер лучше, а каждого пользователя — хуже.

Это не bug, который надо настроить до исчезновения; это сам trade-off, и у каждой стороны есть имя. Latency — то, что ощущает человек, ожидающий ответ. Throughput — то, на что делится счёт. Нет настройки, которая улучшает оба.

Обратите внимание, где это останавливается. От 16 до 32 throughput растёт на 9 %, а latency почти удваивается: шаг перестал быть memory-bound и стал compute-bound, и после этого knee batch ничего не покупает. У каждого deployment есть такой knee; его положение нужно измерять на вашем, но его существование — нет.

Наивный способ batch — собрать $B$ запросов, запустить их вместе и вернуть, когда все закончат. Но они не заканчивают вместе: одни ответы — двадцать token, другие — пятьсот. Fixed batch работает, пока не закончит самый длинный участник, а каждый завершённый запрос до тех пор продолжает занимать свой slot и добавлять padding.

Возьмём 64 запроса с реалистичным перекосом длин output — median 18 token, longest 231, всего 1 874 — и симулируем обе политики при измеренной per-step cost для восьми slots:

Throughput улучшается в 1,6x. Mean latency улучшается более чем в десять раз, потому что при static batching запрос, закончивший за четыре шага, всё равно ждёт соседа на 231 token, прежде чем кто-либо услышит ответ.

Continuous batching³ — исправление, и оно именно настолько просто, как звучит: batch — это не группа, а набор slots, и slot, который освободился, принимает следующий запрос из очереди на следующем же шаге. Scheduler работает с granularностью одного token, а не одного request. Сейчас это делает каждый production serving stack.

У него есть вторая половина — cache. Slots, которые приходят и уходят, оставляют cache memory fragmented, а reservation максимального possible context для каждого slot тратит большую часть reservation впустую. PagedAttention⁴ заимствует ответ у operating systems: хранить cache в fixed-size blocks с block table для каждой sequence, чтобы cache sequence мог быть физически разбросан, оставаясь логически contiguous, — что также позволяет двум sequences с shared prefix делить blocks, которые его хранят. На этом построен vLLM, и поэтому serving engine — это memory allocator с приделанным transformer.

Другая половина счёта — сами weights. Полмиллиарда parameters по четыре bytes каждый — 1,98 GB; по два bytes — 0,99 GB; по одному byte — 0,49 GB. Меньше bits на weight уменьшает model на disk, уменьшает его в memory и — поскольку decode bandwidth-bound — делает каждый шаг быстрее, потому что нужно перемещать меньше bytes.

Самая простая схема — symmetric absolute-maximum quantization, и она помещается в три строки:

qmax  = 2 ** (bits - 1) - 1
scale = W.abs().max() / qmax                        
Wq    = torch.round(W / scale).clamp(-qmax - 1, qmax)
W_hat = Wq * scale                                  # dequantized

Выберите scale так, чтобы самый большой weight отображался в самое большое integer, разделите, округлите, сохраните integers и scale. Восстановите обратным умножением. В этом нет ничего хитрого, и это работает — ровно до момента, когда перестаёт.

Измерено на настоящих weights model: все 168 projection matrices, 357,8 million parameters, relative error $\lVert W - \hat{W}\rVert / \lVert W \rVert$:

Четвёртая строка — это обвал. Relative error 0,99 на worst matrix означает, что reconstruction фактически не сохраняет ничего от original — matrix заменена шумом примерно правильной magnitude. Причина видна в том же experiment на одной matrix:

model.layers.12.mlp.down_proj.weight   (896 x 4864)
mean |w| 0.01386   std 0.01822   max |w| 0.43945   max/std 24.1
weights beyond 6 sigma: 692 of 4,358,144   (0.016 %)

Один weight из шести тысяч лежит дальше шести standard deviations, а largest — на 24. При одном scale для всей matrix этот один weight задаёт step size для всех 4,3 million остальных. При 8 bits есть 256 steps, и typical weight всё ещё попадает на meaningful one. При 4 bits их 16, крайний зарезервирован для значения, которого почти ни у чего нет, а ordinary weights — то есть почти все — округляются до двух-трёх distinct levels.

Всё после этой строки — один и тот же ремонт с разной granularity: дать scale меньшую территорию. Per output row делит error на 3,4; per group of 128 consecutive weights делит его снова. Цена — bookkeeping: 16-bit scale на group of 128 — это $4 + 16/128 = 4.125$ bits на weight вместо 4 — и это возвращает большую часть разрыва.

NF4 подходит с другой стороны.⁵ Levels не обязаны быть equally spaced. Weights внутри block приблизительно normally distributed, поэтому выберите шестнадцать levels как quantiles of a normal distribution: плотные near zero, где weights действительно находятся, редкие в tails, где их нет. Те же four bits, тот же block scaling, но меньший block — 4,25 bits на weight против 4,125 у group-128 — и измеренный error падает с 0,1323 до 0,0952, на 28 % ниже. Часть этого — более fine block, остальное — размещение levels там, где находится mass, а чтобы разделить эти два фактора, понадобилась бы третья строка.

Блок о floating-point в главе 2 закончился обещанием: эта глава quantize weights до 8 и 4 bits и найдёт handful of outlier features, которые откажутся сжиматься. Вот они, и они объясняют, почему «просто округлить числа» никогда не должно было сработать на activations.

Weights выше вели себя плохо. Activations — совершенно другой класс. Возьмите обычный 84-token prompt, capture residual stream на каждом layer и измерьте largest magnitude, которой достигает каждая из 896 dimensions:

Dimension 62 достигает 1 579,6, тогда как median dimension никогда не превышает 1,6. Это не случайность одного token или одного layer: та же dimension есть на layer 4 и всё ещё есть на layer 20, почти с тем же value. Это outlier features,⁶ и они систематичны — свойство trained model, а не input.

Histogram этих 896 per-dimension maxima на layer 16 делает форму безошибочной:

     0 -      1 | ######################################## 254
     1 -      2 | ######################################## 283
     2 -      4 | ######################################## 226
     4 -      8 | ######################################## 93
     8 -     16 | ##################                       18
    16 -     32 | #########                                9
    32 -     64 | #######                                  7
    64 -    128 | #####                                    5
   128 -    256 |                                          0
   256 -    512 |                                          0
   512 -   1024 |                                          0
  1024 -   4096 | #                                        1

Девятьсот dimensions аккуратной кучей ниже 8, затем три октавы вообще ничего, и потом одна dimension в одиночестве на дальнем конце. Теперь quantize этот tensor в INT8 и посчитайте, что происходит:

Четырнадцать levels из 256. Scale был задан 1 579,6, поэтому каждый step шириной 12,44, и typical activation — median magnitude 0,26, ninety-ninth percentile 2,51 — некуда приземлиться. Per dimension это ещё резче:

single tensor-wide scale = 12.4378
  dim 826 (max |h| = 4.77):  1 distinct level out of 256
  dim 336 (max |h| = 1.62):  1 distinct level out of 256
  dim  96 (max |h| = 0.69):  1 distinct level out of 256

after excluding the top 4 dimensions, scale = 0.5749  (22x smaller)
  dim 826: 8 levels    dim 336: 4 levels    dim  96: 3 levels

Один level. Вся dimension, каждый token, quantized в одно и то же число. Восемь bits были выделены, но использовано примерно ноль, и model, читающий эти activations, получает constant.

Это измерение — обоснование для каждой technique, которую люди действительно используют:

Почти каждая статья о quantization останавливается на предыдущем разделе: объясняет method, цитирует compression ratio и утверждает, что quality «в основном сохраняется». Глава 4 была о том, как не обманывать себя, так что выясним.

Тот же model, weights quantized in place каждой scheme, затем три измерения: perplexity на 2 048 token held-out English prose — здесь это draft данного курса, поэтому repository подставляет fixed public-domain book и печатает table той же формы с другими numbers, — battery из 16 short factual questions с known answers при greedy decoding и fraction token, на которых quantized model соглашается с full-precision one при identical context.

В этой таблице четыре вещи стоит сказать прямо.

Правильно сделанный INT8 бесплатен. Per-row INT8 даёт 22,96 против 23,08 у reference — разрыв в одну двухсотую, то есть noise, который следует читать как «identical». Направление noise нестабильно: на public-domain corpus из repository те же две schemes дают 22,24 против 22,18 — половина этой дистанции и в другую сторону. Он соглашается с full-precision model на 142 из 144 generated token. Четверть памяти против fp32 reference, половина против fp16, который вы реально deploy, и никакой detectible cost. INT8, сделанный небрежно, тоже почти бесплатен: один scale на matrix стоит 0,5 points perplexity и ни одного ответа battery. Восемь bits достаточно forgiving, чтобы granularity почти не имела значения, и именно поэтому люди обобщают с INT8 на INT4 и страдают.

INT4 с одним scale на tensor уничтожает model. Perplexity 365 million: не degraded, а annihilated. Затем granularity — вся игра: per-tensor 365 416 000, per-row 46,18, per-group-of-128 31,08, NF4 24,55. Те же four bits на weight, factor fifteen million между worst и best.

Perplexity — грубый инструмент, а battery ещё грубее. Между NF4 и group-128 INT4 gap perplexity равен 6,5 points, а battery отличается на один question — и confidence interval из главы 4 говорит, что один question из sixteen не различает вообще ничего. Есть демонстрация острее interval: запустите тот же battery с выключенным stock repetition penalty model — что greedy decoding actually means, — и эти две строки поменяются местами. Один question из sixteen — не small effect, а no effect. Предупреждение главы 8 тоже применимо: perplexity сопоставима только между models с общим tokenizer, поэтому число из чужого write-up нельзя сравнивать с вашим.

Столбец agreement — самый острый из трёх, и почти бесплатный: запустите full-precision model greedily, затем спросите quantized one в каждой position, что он выбрал бы при том же prefix. У него 144 independent observations вместо 16, ему не нужна ground truth, и он деградирует плавно там, где battery деградирует скачками. Это также ровно та величина, которая нужна следующему разделу.

Это обещание, которое глава 1 дала об этой главе, приходит по расписанию: mathematics говорит, что 4-bit model возможен, а engineering решает, usable ли он.

Глава 12 объявила это и оставила счёт здесь.

Идея следует напрямую из разделения prefill/decode. Проверка предложенной sequence из $\gamma$ token стоит один forward pass по $\gamma$ positions — matrix-matrix product, едва дороже pass по одному. Поэтому:

Output distribution не меняется. При greedy decoding это очевидно — token принимается только если target произвёл бы его. При sampling нужна modified acceptance rule, и Leviathan et al. доказывают, что resulting distribution ровно target distribution.¹⁰ Это вторая exact optimisation в этой главе.

Следовательно, всё зависит от acceptance rate $\alpha$, которую можно измерить, — это agreement column выше, поэтому он там и был посчитан. Используя каждый quantized model как draft для full-precision target, на 144 generated positions:

Expected tokens accepted за verification pass при draft length $\gamma$ равен

а net speedup делит это на собственную стоимость draft, fraction $c$ от target на token:

Bold entry — то, что нужно запомнить: speculative decoding может сделать generation медленнее. При acceptance 30 % и draft, который стоит fifth of the target, вы платите за five forward passes и сохраняете 1,4 token. Последний столбец — другая ловушка: longer draft помогает только при high acceptance, потому что tail of a $\gamma$-token guess почти никогда не достигается. При 90 % acceptance $\gamma = 8$ стоит 3,40x, а при 30 % — 0,79x: одна и та же configuration, win или loss в зависимости от числа, измеренного на вашем traffic.

Quantization уменьшает model, сохраняя ту же function в меньшем числе bits. Distillation уменьшает его, обучая smaller model имитировать larger one¹¹ — идея, появившаяся почти за decade до deep learning.¹²

Тонкость в том, на чём учится student. Не на правильном ответе: на нём его можно было обучить напрямую. То, что добавляет teacher, — whole distribution. Спросите model, что следует за phrase, и посмотрите дальше argmax:

"She poured the milk into the"
  ' jug' 0.1355   ' cup' 0.1051   ' bowl' 0.0605   ' large' 0.0380   ' milk' 0.0360

Hard label говорит jug и больше ничего. Soft label говорит jug, а также что cup было почти так же хорошо, bowl plausible, и large — adjective, совершенно другое grammatical continuation — всё ещё живо. Это исходный аргумент: это 7, но выглядит довольно похоже на 1, и resemblance — information, которую hard label выбрасывает.

Поэтому distillation использует temperature. Деление logits на $T$ перед softmax выравнивает distribution и повышает relative weight runners-up: на этой phrase ratio между top token и third падает с 2,24 при $T = 1$ до 1,50 при $T = 2$ — square root of the first, что и делает деление logits на два с ratio. Тот же ordering, больше attention loss на near misses. Gradient student несёт uncertainty teacher, а не только его verdict.

Всё в этой главе теперь — одна сумма:

где $T$ — общее число resident token по всем concurrent requests. Применим её: строки 7B и 70B предполагают 8 key-value heads размерности 128, строка 13B — полный multi-head attention с 40 heads, как были устроены те поколения model, — и это заметно.

8 GB

16 GB

24 GB

Посмотрите на строку 13B в таблице 8 GB. Weights помещаются — 6,2 GB из 8, — поэтому в обычной манере речи 13B model «работает на карте 8 GB». У него 337 token context, что не conversation, а едва prompt. «Помещается ли он» — неправильный вопрос. Правильный: «с каким context и для скольких пользователей одновременно».

Посмотрите также на две строки int8 в 16 GB. 7B получает 65 378 token, а 13B — 3 136: двадцатикратная разница из-за 5,6 GB дополнительных weights, потому что у этого 13B multi-head attention и его cache стоит 800 KB на token против 128 KB у 7B. Два model похожего size, один unusable для long context, по причине, которая не появляется в headline ни одной model card.

Тринадцать глав назад это был perceptron с двумя weights и bias. Теперь это transformer, который был designed, trained, aligned, научен тратить compute на сложные вопросы и served с измеренной стоимостью на token — и в нём не осталось ни одного unopened box.

Здесь это заканчивается, и заканчивается намеренно.

Глава 14 начинается с model, находящегося где-то ещё. Не в вашем process, не в вашей memory, не в variable, который можно print: на машине, которую вы не администрируете, за API key, port и bill. Всё измеренное здесь всё ещё происходит — prefill всё ещё выполняется перед first token, cache всё ещё растёт вместе с conversation, batch, в котором вы находитесь, всё ещё принадлежит кому-то другому и всё ещё решает вашу latency, — но теперь вы наблюдаете это через stream of Server-Sent Events, finish_reason и HTTP 429 с header Retry-After. Вопросы меняются вместе с vantage point: не как вычисляется этот gradient, а почему мой счёт утроился. Меняется и language, и глава 14 объясняет это правило, а не объявляет его: до этого code держал weights, gradients, logits и tokenizer bytes; дальше он держит connection, retry, cancellation и accumulated state. Тринадцать глав позади не отбрасываются при переходе. Они — описание того, что работает по другую сторону port.

Два omission намеренны. FlashAttention (Dao et al., arXiv:2205.14135) — не другой attention: он вычисляет ту же function, разбивая operation на tiles так, чтобы score matrix $n \times n$ никогда не записывалась в memory, поэтому 67 MB во второй таблице этой главы на практике меньше, чем suggests arithmetic. А сами kernels delegated: lecture 10 Stanford CS336 покрывает inference systems в глубине, на которую эта глава не претендует, а repository llama.cpp и specification GGUF — primary sources для CPU side.