Как да направим inference евтин: KV cache, batching и квантизация

За да генерира token, decoder-only transformer взема цялата досегашна последователност, прекарва я през всеки слой и прочита вероятностното разпределение от последната позиция. След това добавя избрания token и прави същото отново. Това описание е вярно и точно това прави бавното изпълнение.

То е и изключително разточително, а причината е causal mask от Глава 9. Векторите key и value на позиция 7 се изчисляват от входа на позиция 7 и позициите преди нея. Когато пристигне позиция 8, позиция 7 не може да я вижда — това означава causal — така че key и value на позиция 7 са точно същите числа като преди. Бавното изпълнение въпреки това ги преизчислява, на всяка стъпка.

Затова ги съхранете. Това хранилище е key-value cache, най-важната оптимизация в serving на езикови модели:

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 всеки кеширан key, а кешираните keys така или иначе нямаше да се променят. Това не е приближение — проверката за идентичен изход по-горе е смисълът. Cache не разменя качество срещу скорост; той изтрива излишна аритметика.

За да се види мащабирането чисто, махнете transformer и измерете една attention head с $d = 64$, една стъпка на генериране, изчислена и по двата начина:

Дясната колона е причината. Преизчисляването изгражда пълната $n \times n$ attention матрица на всяка стъпка — $O(n^2)$ от карето с асимптотична нотация в Глава 9, платено веднъж за token. С cache вместо това изграждате ред $1 \times n$: при 4 096 tokens — 67 MB scores срещу 16 KB.

Броенето на multiply-accumulates вместо милисекунди премахва машината от аргумента. За генериране на $T$ tokens от студен старт:

На стъпка кешираната версия е линейна спрямо context, а некешираната — квадратична; сумирано върху генериране, $O(T^2)$ срещу $O(T^3)$, като съотношението расте без ограничение. Деветкратната разлика в началото беше измерена върху 48 tokens — под първия ред на тази таблица.

Cache променя и това, което трябва да бъде в паметта. На лаптоп GPU с 8 GB, генериращ 256 tokens във fp16, като вземем пика на allocator и извадим resident weights:

8,3 пъти повече памет, похарчена, за да произведе същите tokens по-бавно. Това е обещанието от Глава 5, пристигащо от неочаквана посока: там reverse-mode autodiff трябваше да държи всеки междинен резултат жив за backward pass, а активациите доминираха паметта при обучение. При 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, същият хардуер, същите тегла, шесткратна разлика на token — и посоката е обратна на това, което повечето хора очакват. Prompt е евтината част. Генерирането се разделя на две фази с наистина различна физика:

И двете половини са измерими. Prefill, един pass върху $P$ tokens:

Decode, един token срещу cache от $C$:

Прочетете втората таблица два пъти. Преминаването от 16 tokens context към 1 024 — шестдесет и четири пъти повече история, върху която да се attend — не промени цената на стъпката с нищо измеримо. Attention срещу cache е реална работа, но е засенчена от фиксираната цена да прекарате половин милиард тегла през шината на паметта, за да произведете един вектор. Тази фиксирана цена е причината за всичко в следващия раздел.

Тези две фази са произходът на двете числа, които всяка serving система отчита. Time to first token е по същество prefill и расте с prompt, затова дълъг разговор се усеща бавен в началото. Tokens per second е $1/\text{decode step}$ и е приблизително постоянен, затова отговорът после тече равномерно. Chat, който започва бавно и след това streams плавно, не е трик на визуализацията. Това са тези две таблици.

Cache разменя аритметика срещу памет, а паметта, която иска, не е малка. За всеки token в context всеки слой държи един key вектор и един value вектор за всяка key-value head:

2 е за keys и values; всичко останало е архитектурата. За model, измерван през тази глава — 24 слоя, 14 query heads, 2 key-value heads, head dimension 64 — във fp16 това е $2 \times 24 \times 2 \times 64 \times 2 = 12{,}288$ bytes на token.

Формулите в тази област имат навика да грешат с фактор две, така че проверете я срещу 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

Точно, и остава точно при всяка изпробвана форма:

Последните три реда заслужават втори поглед. Тридесет и двама потребители с по 2 048 tokens, шестдесет и четирима с по 1 024, сто двадесет и осем с по 512 — cache е 768 MB във всеки случай, защото и трите държат 65 536 tokens. Cache зависи само от общия брой resident tokens, не от това как са разпределени между потребителите. Този факт е основата на раздела за 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 ограничаващото условие.

Което става бързо. За model от 7B клас с 32 слоя и 8 key-value heads с размерност 128, cache е 128 KB на token във fp16:

Собствените тегла на този model са 13,0 GB във fp16, числото в таблицата в края на тази глава. Така че при context от 128 000 tokens cache на един потребител е по-голям от model. Това е аритметиката, която Глава 16 превръща в пари, и затова дълъг разговор не е просто бавен — той заема фиксиран дял от машина, докато заявката е жива.

Decode е memory-bound: теглата се влачат през шината, за да произведат един token, а аритметичните блокове бездействат. Затова сложете повече работа в същата стъпка. Пуснете няколко заявки едновременно и теглата, прочетени веднъж, обслужват всички тях. Измерено на същия model, като всяка заявка държи 64-token cache и декодира един token:

Прочетете двете десни колони една срещу друга, защото те са целият смисъл. Преминаването от една заявка към шестнадесет умножава throughput по 6,0 и умножава чакането за всяка отделна заявка по 2,67. Batch направи сървъра по-добър и всеки потребител по-зле.

Това не е bug, който да се настрои away; това е самият компромис, и има име от всяка страна. Латентност е това, което изпитва човек, чакащ отговор. Throughput е това, на което се дели фактурата. Няма настройка, която да подобрява и двете.

Обърнете внимание и къде спира. От 16 до 32 throughput печели 9 %, докато латентността почти се удвоява: стъпката е спряла да бъде memory-bound и е станала compute-bound, а след това коляно batch не купува нищо. Всяко deployment има такова коляно; местоположението му трябва да се измери на вашето, но съществуването му — не.

Наивният начин за batching е да съберете $B$ заявки, да ги пуснете заедно и да върнете резултат, когато всички приключат. Но те не приключват заедно: някои отговори са двадесет tokens, а други — петстотин. Фиксиран batch работи, докато най-дългият му член приключи, и всяка завършена заявка продължава да заема слота си, добавяйки padding, дотогава.

Вземете 64 заявки с реалистично изкривяване на дължините на изхода — медиана 18 tokens, най-дълга 231, общо 1 874 — и симулирайте двете политики при измерената цена на стъпка за осем слота:

Throughput се подобрява 1,6x. Средната латентност се подобрява повече от десет пъти, защото при static batching заявка, приключила за четири стъпки, все още чака 231-token съсед, преди някой да чуе за нея.

Continuous batching³ е поправката и е толкова проста, колкото звучи: batch не е група, а набор от слотове, и слот, който се освободи, допуска следващата чакаща заявка още на следващата стъпка. Scheduler работи с гранулярност един token, а не една заявка. Всяка serving stack в production вече прави това.

Има и втора половина — cache. Слотовете, които идват и си отиват, оставят cache паметта фрагментирана, а резервирането на максималния възможен context за всеки слот губи по-голямата част от резервацията. PagedAttention⁴ заема отговора от операционните системи: съхранявайте cache във фиксирани по размер блокове с block table за всяка последователност, така че cache на последователността да може да бъде физически разпръснат, докато остава логически непрекъснат — което също позволява две последователности със споделен prefix да споделят блоковете, които го държат. На това е изграден vLLM и затова serving engine е memory allocator с прикрепен transformer.

Другата половина от сметката са самите тегла. Половин милиард параметъра по четири bytes са 1,98 GB; по два bytes — 0,99 GB; по един byte — 0,49 GB. По-малко bits на тегло свиват model на диска, свиват го в паметта и — понеже 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, така че най-голямото тегло да се мапва към най-голямото цяло число, разделете, закръглете, съхранете целите числа и scale. Възстановете чрез умножаване обратно. В нея няма нищо хитро и работи — чак докато спре.

Измерено върху реалните тегла на model: всички 168 projection matrices, 357,8 милиона параметъра, относителна грешка $\lVert W - \hat{W}\rVert / \lVert W \rVert$:

Четвъртият ред е сривът. Относителна грешка 0,99 на най-лошата матрица означава, че реконструкцията не запазва по същество нищо от оригинала — матрицата е заменена с шум с приблизително правилната големина. Причината се вижда в същия експеримент върху една матрица:

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 %)

Едно тегло от шест хиляди е отвъд шест стандартни отклонения, а най-голямото е на 24. С един scale за цялата матрица това едно тегло задава размера на стъпката за всичките 4,3 милиона. При 8 bits има 256 стъпки и типичното тегло все още попада на смислена. При 4 bits има 16, най-външната е запазена за стойност, която почти нищо няма, а обикновените тегла — тоест всички — се закръглят до две или три различни нива.

Всичко след този ред е същата поправка с различна гранулярност: дайте на scale по-малка територия. Per output row дели грешката на 3,4; per group от 128 последователни тегла я дели отново. Цената е счетоводство — 16-bit scale на група от 128 е $4 + 16/128 = 4.125$ bits на тегло вместо 4 — и връща по-голямата част от разликата.

NF4 подхожда от другата страна.⁵ Нивата не трябва да са равномерно разположени. Теглата в block са приблизително нормално разпределени, затова изберете шестнадесетте нива като квантилите на нормално разпределение: гъсти близо до нулата, където теглата наистина са, редки в опашките, където не са. Същите четири bits, същото block scaling, при по-малък block — 4,25 bits на тегло срещу 4,125 за group-128 — и измерената грешка пада от 0,1323 до 0,0952, 28 % по-ниско. Част от това е по-финият block, останалото е поставянето на нивата там, където е масата, а разделянето на двете би изисквало трети ред.

Карето за floating-point в Глава 2 завърши с обещание: че тази глава ще квантизира тегла до 8 и 4 bits и ще намери шепа outlier features, отказващи да бъдат притиснати. Ето ги, и те обясняват защо „просто закръглете числата“ никога нямаше да работи върху активации.

Теглата по-горе се държаха лошо. Активациите са в друга лига. Вземете обикновен 84-token prompt, прихванете residual stream на всеки слой и измерете най-голямата големина, която всяко от 896-те измерения достига:

Измерение 62 достига 1 579,6, докато медианното измерение никога не надхвърля 1,6. Това не е случайност на един token или един слой: същото измерение е там на слой 4 и още е там на слой 20, с почти същата стойност. Това са outlier features,⁶ и те са систематични — свойство на обучения model, не на входа.

Хистограмата на тези 896 максимума по измерение на слой 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

Деветстотин измерения в подредена купчина под 8, абсолютно нищо за три октави, после едно измерение само в далечния край. Сега квантизирайте този tensor до INT8 и пребройте какво става:

Четиринадесет нива от 256. Scale беше зададен от 1 579,6, така че всяка стъпка е широка 12,44, а типичната активация — медианна големина 0,26, деветдесет и девети percentile 2,51 — няма къде да попадне. По измерение е още по-рязко:

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

Едно ниво. Цялото измерение, всеки token, квантизирано до едно и също число. Осем bits бяха разпределени и приблизително нула бяха използвани, а model, който чете тези активации, получава константа.

Това измерване е оправданието за всяка техника, която хората реално използват:

Почти всяка статия за квантизация спира в предишния раздел: обяснява метода, цитира compression ratio и твърди, че качеството е „до голяма степен запазено“. Глава 4 беше за това да не се самозалъгвате, така че нека разберем.

Същият model, теглата квантизирани на място с всяка схема, после три измервания: perplexity върху 2 048 tokens задържана английска проза — тук, черновата на този курс, затова repository заменя с фиксирана public-domain книга и отпечатва таблица със същата форма и различни числа — набор от 16 кратки фактически въпроса с известни отговори при greedy decoding, и делът на tokens, при които квантизираният model се съгласява с full-precision model при идентичен context.

Четири неща в тази таблица си струва да бъдат казани ясно.

INT8, направен правилно, е безплатен. Per-row INT8 дава 22,96 срещу 23,08 на reference — разлика от една част на двеста, което е шум и трябва да се чете като „идентично“. Накъде сочи шумът не е стабилно: върху public-domain corpus в repository същите две схеми излизат 22,24 срещу 22,18: половината разстояние, и сочещо в другата посока. Той се съгласява с full-precision model върху 142 от 144 генерирани tokens. Четвърт от паметта спрямо fp32 reference, половина спрямо fp16, който реално бихте deploy-нали, и без откриваема цена. INT8, направен небрежно, също е почти безплатен: един scale на матрица струва 0,5 perplexity points и никакви отговори от набора. Осем bits са достатъчно forgiving, че гранулярността почти няма значение, което е точно причината хората да обобщават от INT8 към INT4 и да пострадат.

INT4 с един scale на tensor унищожава model. Perplexity 365 милиона: не влошен, а анихилиран. Гранулярността тогава е цялата игра — per-tensor 365 416 000, per-row 46,18, per-group-of-128 31,08, NF4 24,55. Същите четири bits на тегло, фактор петнадесет милиона между най-лошото и най-доброто.

Perplexity е груб инструмент, а наборът — още по-груб. Между NF4 и group-128 INT4 разликата в perplexity е 6,5 points, а наборът се различава с един въпрос — и confidence interval от Глава 4 казва, че един въпрос от шестнадесет не различава абсолютно нищо. Има по-остра демонстрация от интервала: пуснете същия набор с изключена стандартна repetition penalty на model, което всъщност означава greedy decoding, и тези два реда си разменят местата. Един въпрос от шестнадесет не е малък ефект, а никакъв ефект. Предупреждението от Глава 8 също важи: perplexity е сравнима само между модели, споделящи tokenizer, така че число от чужд write-up не може да се сравнява с вашето.

Колоната за съгласие е най-острата от трите, и почти безплатна: пуснете full-precision model greedily, после попитайте квантизирания на всяка позиция какво би избрал при същия prefix. Има 144 независими наблюдения вместо 16, не изисква ground truth и се влошава плавно там, където наборът се влошава на скокове. Тя е и точно величината, от която се нуждае следващият раздел.

Това е обещанието, което Глава 1 даде за тази глава, пристигащо по график: математиката казва, че 4-bit model е възможен, а инженерството решава дали е използваем.

Глава 12 обяви това и остави сметката тук.

Идеята идва директно от разделението prefill/decode. Проверяването на предложена последователност от $\gamma$ tokens струва един forward pass върху $\gamma$ позиции — matrix-matrix product, едва по-скъп от pass върху една. Така че:

Изходното разпределение е непроменено. При greedy decoding това е очевидно — token се приема само ако target би го произвел. При sampling изисква модифицирано правило за приемане, а Leviathan et al. доказват, че полученото разпределение е точно target.¹⁰ Това е втората точна оптимизация в тази глава.

Следователно всичко зависи от acceptance rate $\alpha$, който е измерим — това е колоната за съгласие по-горе, затова беше изчислена там. Използвайки всеки квантизиран model като draft за full-precision target, върху 144 генерирани позиции:

Очакваните tokens, приети на verification pass, при draft length $\gamma$, са

а нетното ускорение дели това на собствената цена на draft, дял $c$ от target на token:

Удебеленият запис е този, който трябва да запомните: speculative decoding може да направи генерирането по-бавно. При 30 % acceptance с draft, струващ една пета от target, плащате за пет forward passes и запазвате 1,4 tokens. Последната колона е другият капан — по-дълъг draft помага само когато acceptance е висока, защото опашката на $\gamma$-token предположение почти никога не се достига. При 90 % acceptance $\gamma = 8$ струва 3,40x, а при 30 % струва 0,79x: същата конфигурация, печалба или загуба според число, измерено върху вашия трафик.

Квантизацията свива model, като съхранява същата функция в по-малко bits. Дистилацията го свива, като обучава по-малък model да имитира по-голям¹¹ — идея, която предхожда deep learning с почти десетилетие.¹²

Фината част е от какво учи student. Не от правилния отговор: можеше да бъде обучен директно върху него. Това, което teacher добавя, е цялото разпределение. Попитайте model какво следва след фраза и погледнете отвъд 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 — правдоподобно, и large — прилагателно, напълно различно граматично продължение — все още живо. Това е оригиналният аргумент: това е 7, но прилича доста на 1, а приликата е информация, която hard label изхвърля.

Затова дистилацията използва и температура. Делението на logits на $T$ преди softmax изглажда разпределението и повишава относителната тежест на подгласниците: при тази фраза съотношението между top token и третия пада от 2,24 при $T = 1$ до 1,50 при $T = 2$ — квадратният корен на първото, което е ефектът от делене на logits на две върху съотношение. Същият ред, повече от attention на loss върху близките пропуски. Gradient на student носи несигурността на teacher, а не само присъдата му.

Всичко в тази глава вече е една сума:

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

8 GB

16 GB

24 GB

Погледнете реда 13B в таблицата за 8 GB. Теглата се побират — 6,2 GB от 8 — така че по обичайния начин на говорене 13B model „работи на 8 GB карта“. Той има 337 tokens context, което не е разговор, а едва prompt. „Побира ли се“ е грешният въпрос. Правилният е „с колко context и за колко потребители едновременно“.

Погледнете и двата int8 реда при 16 GB. 7B получава 65 378 tokens, а 13B — 3 136 — двадесеткратна разлика от 5,6 GB допълнителни тегла, защото 13B тук има multi-head attention и неговият cache струва 800 KB на token срещу 128 KB на 7B. Два model с подобен размер, единият неизползваем за дълъг context, по причина, която не се появява в заглавието на никоя model card.

Преди тринадесет глави това беше perceptron с две тегла и bias. Сега е transformer, който е проектиран, обучен, aligned, научен да харчи compute за трудни въпроси и served при измерена цена на token — без нито една неотворена кутия вътре.

Това приключва тук, и приключва нарочно.

Глава 14 започва с model някъде другаде. Не във вашия process, не във вашата памет, не в променлива, която можете да print-нете: на машина, която не администрирате, зад API key, port и bill. Всичко измерено тук все още се случва — prefill все още върви преди първия token, cache все още расте с разговора, batch, в който сте, все още принадлежи на някой друг и все още решава вашата латентност — но оттук нататък го наблюдавате през stream от Server-Sent Events, finish_reason и HTTP 429 с header Retry-After. Въпросите се променят с гледната точка: не как се изчислява този gradient, а защо фактурата ми стана тройна. Езикът също, и Глава 14 обяснява това правило, вместо да го обявява — дотук code държеше weights, gradients, logits и tokenizer bytes; оттам нататък държи connection, retry, cancellation и accumulated state. Тринадесетте глави зад вас не се изхвърлят при преминаването. Те са описанието на това, което работи от другата страна на port.

Два пропуска са умишлени. FlashAttention (Dao et al., arXiv:2205.14135) не е различен attention — той изчислява същата функция чрез tiling на операцията, така че $n \times n$ score matrix никога да не се записва в паметта, затова 67 MB във втората таблица на тази глава са по-малки на практика, отколкото аритметиката предполага. А самите kernels са делегирани: лекция 10 от Stanford CS336 покрива inference systems в дълбочина, която това не се опитва да постигне, а repository llama.cpp и спецификацията GGUF са основните източници за CPU страната.