Як здешевити inference: KV cache, batching і квантування

Щоб згенерувати token, decoder-only transformer бере всю sequence на цей момент, проганяє її через кожен layer і зчитує probability distribution з останньої позиції. Потім він додає вибраний token і робить це знову. Цей опис правильний, і саме це робить повільний запуск.

Але це ще й надзвичайно марнотратно, і причина — causal mask із Розділу 9. Key і value vectors позиції 7 обчислюються з input позиції 7 і позицій перед нею. Коли приходить позиція 8, позиція 7 не може її бачити — саме це означає causal — тож key і value позиції 7 є точно тими самими числами, що й раніше. Повільний запуск усе одно переобчислює їх на кожному кроці.

Тож збережіть їх. Це сховище — key-value cache, найважливіша optimization у serving мовних models:

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 усередині loop: nxt, один token. Не sequence. Query нового token attends до кожного cached key, а cached keys і так ніколи не мали змінитися. Це не approximation — перевірка ідентичності output вище і є суттю. Cache не міняє якість на швидкість; він видаляє надлишкову arithmetic.

Щоб чисто побачити 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 tokens це 67 MB scores проти 16 KB.

Підрахунок multiply-accumulates замість milliseconds прибирає машину з аргументу. Щоб згенерувати $T$ tokens з cold start:

На кожному кроці cached version є linear щодо context, а uncached — quadratic; у сумі за generation це $O(T^2)$ проти $O(T^3)$, причому ratio росте без межі. Дев’ятикратну різницю на початку було виміряно на 48 tokens — ще до першого рядка тієї таблиці.

Cache також змінює те, що має бути в memory. На 8 GB laptop GPU, який генерує 256 tokens у fp16, якщо взяти peak allocator і відняти resident weights:

У 8.3 раза більше memory, витраченої на те, щоб створити ті самі tokens повільніше. Це обіцянка з Розділу 5, яка приходить із несподіваного боку: там reverse-mode autodiff мав тримати кожен intermediate живим для backward pass, і activations домінували training memory. В inference немає backward pass і нічого не треба зберігати для нього — тож memory натомість домінує 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, а кожен generated token — 166 ms. Той самий model, те саме hardware, ті самі weights, шестикратна різниця за token — і вона спрямована не туди, куди очікує більшість. Prompt — це дешева частина. Generation ділиться на дві фази з по-справжньому різною фізикою:

Обидві половини можна виміряти. Prefill, один pass через $P$ tokens:

Decode, один token проти cache з $C$:

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

Ці дві фази — джерело двох чисел, які показує кожна serving system. Time to first token — це по суті prefill, і він росте разом із prompt, тому довга розмова повільно стартує. Tokens per second — це $1/\text{decode step}$, і він приблизно сталий, тому відповідь далі тече рівномірно. Chat, який повільно починається, а потім плавно streaming, — це не rendering trick. Це дві наведені таблиці.

Cache міняє arithmetic на memory, і memory, якої він хоче, не мала. Для кожного 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.

Formulae в цій сфері мають звичку помилятися вдвічі, тож перевірте це на 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:

Останні три рядки заслуговують на другий погляд. Тридцять два users із 2,048 tokens кожен, шістдесят чотири з 1,024, сто двадцять вісім з 512 — cache у кожному випадку 768 MB, бо всі три варіанти тримають 65,536 tokens. Cache залежить лише від загальної кількості resident tokens, а не від того, як вони розподілені між users. Цей факт — фундамент розділу про batching.

Розділ 9 представив multi-query і grouped-query attention, відклавши причину до цього розділу. Причина — ця formula, а саме $H_{kv}$ у ній.

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

Multi-query attention¹ доводить це до межі: усі query heads ділять одну key-value head. Grouped-query attention² — компроміс, який переміг: кілька key-value heads, кожна з яких спільна для групи query heads, — бо втрата якості MQA була реальною, а в GQA її немає. Жоден із них не купує arithmetic. Вони існують, щоб поділити цю formula на integer, і поширилися індустрією тієї миті, коли long contexts зробили cache binding constraint.

А він робить це швидко. Для 7B-class model із 32 layers і 8 key-value heads dimension 128 cache становить 128 KB на token у fp16:

Власні weights цього model — 13.0 GB у fp16, цифра в таблиці наприкінці розділу. Тож при context у 128,000 tokens cache одного user більший за model. Це arithmetic, яку Розділ 16 перетворює на гроші, і тому довга розмова не просто повільна — вона займає фіксований шматок машини весь час, поки request живий.

Decode є memory-bound: weights протягуються через bus, щоб створити один token, а arithmetic units простоюють. Тож покладіть більше роботи в той самий крок. Запустіть кілька requests одночасно, і weights, прочитані один раз, обслужать їх усі. Виміряно на тому самому model, кожен request тримає 64-token cache і decode один token:

Читайте два праві стовпці один проти одного, бо в них уся суть. Перехід від одного request до шістнадцяти множить throughput на 6.0 і множить очікування будь-якого окремого request на 2.67. Batch зробив server кращим, а кожного user — гіршим.

Це не bug, який можна налаштуванням прибрати; це сама tradeoff, і він має назву з кожного боку. Latency — це те, що відчуває людина, яка чекає відповіді. Throughput — це те, на що ділиться invoice. Жоден setting не покращує обидва.

Зверніть увагу й на те, де це зупиняється. Від 16 до 32 throughput зростає на 9 %, а latency майже подвоюється: крок перестав бути memory-bound і став compute-bound, а після цього knee batch уже нічого не купує. Кожен deployment має такий knee; його місце треба виміряти на вашому, але його існування — ні.

Наївний спосіб batch — зібрати $B$ requests, запустити їх разом і повернути результат, коли всі завершаться. Але вони не завершуються разом: деякі replies мають двадцять tokens, а деякі — п’ятсот. Fixed batch працює, доки не завершиться найдовший member, і кожен завершений request до того часу продовжує займати свій slot, додаючи padding.

Візьміть 64 requests із реалістично скошеними output lengths — median 18 tokens, longest 231, 1,874 загалом — і змоделюйте обидві policies за виміряною per-step cost для восьми slots:

Throughput поліпшується в 1.6x. Mean latency поліпшується більше ніж у десять разів, бо за static batching request, який завершився за чотири steps, усе одно чекає на 231-token сусіда, перш ніж хтось почує результат.

Continuous batching³ — це виправлення, і воно настільки просте, як звучить: batch — це не група, а набір slots, і slot, який звільнився, приймає наступний request із queue вже на наступному step. Scheduler працює з granularity одного token, а не одного request. Кожен production serving stack тепер робить саме це.

У нього є друга половина — cache. Slots, які приходять і йдуть, залишають cache memory fragmented, а резервування для кожного slot його максимально можливого context марнує більшість reservation. PagedAttention⁴ позичає відповідь в operating systems: зберігати cache у fixed-size blocks із block table для кожної sequence, щоб cache sequence міг бути фізично розкиданим, але логічно contiguous — і це також дозволяє двом sequences зі спільним 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 — робить кожен step швидшим, адже треба переміщувати менше bytes.

Найпростіша scheme — 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$:

Четвертий рядок — це collapse. Relative error 0.99 на worst matrix означає, що reconstruction по суті не зберігає нічого від original — matrix замінено noise приблизно правильної 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, а найбільша — на 24. З одним scale для всієї matrix ця одна weight задає step size для всіх 4.3 million інших. При 8 bits є 256 steps, і типова weight усе ще потрапляє на meaningful one. При 4 bits їх 16, крайній зарезервовано для value, якого майже ні в кого немає, а ordinary weights — тобто всі вони — округлюються до двох або трьох distinct levels.

Усе після цього рядка — той самий ремонт на різних granularities: дати 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, тож виберіть sixteen levels як quantiles normal distribution: щільно біля zero, де weights справді є, рідко в tails, де їх немає. Ті самі four bits, той самий block scaling, на меншому block — 4.25 bits на weight проти 4.125 у group-128 — і measured error падає з 0.1323 до 0.0952, на 28 % нижче. Частина цього — finer block, решта — розміщення levels там, де mass, а щоб розділити ці дві причини, потрібен був би третій рядок.

Блок про floating-point у Розділі 2 завершився обіцянкою: цей розділ quantize weights до 8 і 4 bits і знайде жменьку outlier features, які відмовляються стискатися. Ось вони, і вони пояснюють, чому «просто округлити числа» ніколи не мало спрацювати для activations.

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

Dimension 62 досягає 1,579.6, тоді як median dimension ніколи не перевищує 1.6. Це не випадковість одного token чи одного layer: та сама dimension є на layer 4 і все ще є на layer 20, майже з тим самим value. Це outlier features,⁶ і вони systematic — властивість 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, взагалі нічого на три octaves, а потім одна dimension самотньо в далекому кінці. Тепер quantize цей tensor до INT8 і порахуйте, що станеться:

Чотирнадцять levels із 256. Scale задавався 1,579.6, тож кожен step має ширину 12.44, а типовій 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 до того самого number. Було виділено eight bits, а використано приблизно нуль, і model, який читає ці activations, отримує constant.

Це вимірювання — виправдання кожної technique, яку люди реально використовують:

Майже кожна стаття про quantization зупиняється на попередньому розділі: пояснює method, цитує compression ratio і стверджує, що quality «здебільшого збережена». Розділ 4 був про те, як не обдурити себе, тож з’ясуймо.

Той самий model, weights quantized in place кожною scheme, потім три measurements: perplexity на 2,048 tokens відкладеної англійської prose — тут, чернетці цього course, саме тому repository підставляє fixed public-domain book і друкує таблицю тієї самої форми з іншими numbers — battery з 16 short factual questions із known answers під greedy decoding, і fraction of tokens, на яких quantized model погоджується з full-precision model за 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 tokens. Чверть memory проти fp32 reference, половина проти fp16, який ви реально deploy, і жодної detectable cost. INT8, зроблений недбало, теж майже безкоштовний: один scale на matrix коштує 0.5 perplexity points і жодної відповіді в battery. Eight 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 у п’ятнадцять мільйонів між 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, і ці два рядки поміняються місцями. Одне question із sixteen — це не малий effect, це no effect. Попередження Розділу 8 теж застосовне: perplexity можна порівнювати лише між models зі спільним tokenizer, тож число з чужого write-up не можна порівнювати з вашим.

Стовпець agreement — найгостріший із трьох, і майже безкоштовний: запустіть full-precision model greedily, а потім на кожній position запитайте quantized model, що він вибрав би за того самого prefix. Він має 144 independent observations замість 16, не потребує ground truth і деградує плавно там, де battery деградує стрибками. Це також рівно та quantity, яка потрібна наступному розділу.

Це обіцянка, яку Розділ 1 дав щодо цього розділу, вчасно виконана: mathematics каже, що 4-bit model можливий, а engineering вирішує, чи придатний він до використання.

Розділ 12 оголосив це й залишив рахунок тут.

Ідея прямо випливає з поділу prefill/decode. Verification proposed sequence з $\gamma$ tokens коштує один forward pass через $\gamma$ positions — matrix-matrix product, ледь дорожчий за pass через одну. Тож:

Output distribution не змінюється. З greedy decoding це очевидно — token accepted лише якщо target його й produced би. З sampling потрібне modified acceptance rule, і Leviathan et al. доводять, що resulting distribution точно дорівнює target.¹⁰ Це друга exact optimisation у цьому розділі.

Отже, все залежить від acceptance rate $\alpha$, яку можна виміряти — це agreement column вище, саме тому його там обчислено. Використовуючи кожен quantized model як draft для full-precision target, на 144 generated positions:

Expected tokens accepted per verification pass, при draft length $\gamma$, дорівнює

а net speedup ділить це на власну cost draft, fraction $c$ від target на token:

Жирний запис — той, який слід запам’ятати: speculative decoding може зробити generation повільнішою. При 30 % acceptance з draft, що коштує п’яту частину target, ви платите за five forward passes і зберігаєте 1.4 tokens. Останній стовпець — інша пастка: longer draft допомагає лише за high acceptance, бо tail $\gamma$-token guess майже ніколи не досягається. При 90 % acceptance $\gamma = 8$ вартий 3.40x, а при 30 % — 0.79x: та сама configuration, win або loss залежно від number, виміряного на вашому traffic.

Quantization стискає model, зберігаючи ту саму function у меншій кількості bits. Distillation стискає його, training менший model імітувати більший¹¹ — ідея, що майже на decade старша за deep learning.¹²

Тонкість у тому, з чого student learns. Не з правильної answer: на ній його можна було train directly. Те, що додає 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 був майже так само good, bowl plausible, а large — adjective, зовсім інше grammatical continuation — усе ще live. Це original argument: це 7, але воно дуже схоже на 1, і ця resemblance — інформація, яку hard label викидає.

Саме тому distillation використовує temperature. Ділення logit на $T$ перед softmax вирівнює distribution і піднімає relative weight runners-up: у цій phrase ratio між top token і third падає з 2.24 при $T = 1$ до 1.50 при $T = 2$ — square root першого, саме це робить ділення logit на два з ratio. Той самий ordering, більше attention loss на near misses. Gradient student несе uncertainty teacher, а не лише його verdict.

Усе в цьому розділі тепер одна сума:

де $T$ — це total resident tokens across all concurrent requests. Застосуймо: рядки 7B і 70B припускають 8 key-value heads dimension 128, рядок 13B — full multi-head attention із 40 heads, саме так були побудовані ті generations of model — і це видно.

8 GB

16 GB

24 GB

Подивіться на рядок 13B у таблиці 8 GB. Weights вміщуються — 6.2 GB із 8 — тож у звичній манері говорити 13B model «runs on an 8 GB card». Він має 337 tokens context, тобто це не розмова, а ледве prompt. «Чи вміщується?» — неправильне питання. Правильне: «з яким context і для скількох users одночасно».

Подивіться також на два рядки int8 у 16 GB. 7B отримує 65,378 tokens, а 13B — 3,136: двадцятикратна різниця від 5.6 GB extra weights, бо 13B тут має multi-head attention, і його cache коштує 800 KB на token проти 128 KB у 7B. Два models подібного size, один непридатний для long context, з причини, яка не з’являється в headline жодної model card.

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

На цьому все закінчується, і закінчується навмисно.

Розділ 14 починається з model десь в іншому місці. Не у вашому process, не у вашій memory, не в variable, яку можна print: на machine, яку ви не administer, за API key, port і bill. Усе, що вимірювалося тут, усе ще відбувається — prefill усе ще runs перед першим token, cache усе ще grows разом із conversation, batch, у якому ви перебуваєте, усе ще належить комусь іншому й усе ще decides вашу latency — але відтепер ви observes це через stream Server-Sent Events, finish_reason і HTTP 429 з header Retry-After. Questions змінюються разом із vantage point: не як обчислюється цей gradient, а чому мій invoice потроївся. Language теж змінюється, і Розділ 14 пояснює це правило, а не оголошує його — до цього моменту code тримав weights, gradients, logits і tokenizer bytes; далі він тримає connection, retry, cancellation і accumulated state. Тринадцять розділів позаду вас не відкидаються під час переходу. Вони — опис того, що працює по той бік port.

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