推論コストを下げる：KV Cache、Batching、量子化

tokenを生成するために、decoder-only transformerはこれまでのsequence全体を取り、すべてのlayerに通し、最後のpositionから確率分布を読み取ります。それから選ばれたtokenを追加し、また同じことをします。この説明は正しく、遅い実行がやっていることでもあります。

同時に、それは途方もなく無駄です。その理由は第9章のcausal maskにあります。position 7のkey vectorとvalue vectorは、position 7のinputとそれ以前のpositionから計算されます。position 8が到着しても、position 7はそれを見ることができません。causalとはそういう意味です。したがってposition 7のkeyとvalueは、前とまったく同じ数値です。遅い実行は、それでも各stepでそれらを再計算します。

だから保存します。その保存領域がkey-value cacheであり、language model servingにおける最も重要なoptimizationです。

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)

loop内でmodelに渡されているものを見てください：nxt、つまり1 tokenです。sequenceではありません。新しいtokenのqueryはcached keyすべてにattendし、cached keyはそもそも変わるはずがありません。これは近似ではありません。上の同一出力チェックがその要点です。cacheは品質を速度と交換しません。冗長な算術を削除します。

scalingをきれいに見るために、transformerを取り払い、$d = 64$のsingle attention headを、生成1 stepについて両方の方法で測ります。

右端の列が原因です。再計算は毎stepで完全な$n \times n$ attention matrixを作ります。第9章の漸近記法のboxに出てきた$O(n^2)$を、tokenごとに1回支払うわけです。cacheを使うと、代わりに$1 \times n$のrowを作るだけです。4,096 tokensでは、scoresが67 MB対16 KBです。

ミリ秒ではなくmultiply-accumulateを数えると、machineを議論から取り除けます。cold startから$T$ tokensを生成する場合：

stepごとには、cached版はcontextに対してlinearで、uncached版はquadraticです。generation全体で合計すると$O(T^2)$対$O(T^3)$になり、その比率は限りなく大きくなります。冒頭の9倍差は48 tokensで測ったものです。この表の最初のrowより短い長さです。

cacheは、memoryに置くべきものも変えます。8 GBのlaptop GPUでfp16、256 tokensを生成し、allocatorのpeakから常駐weightsを差し引くと：

memoryは8.3倍多く、しかも同じtokensをより遅く生成するために使われます。これは第5章での約束が、思いがけない方向から到着したものです。そこではreverse-mode autodiffがbackward passのためにすべての中間値を生かしておく必要があり、activationsがtraining memoryを支配しました。inferenceにはbackward passがなく、そのために保持すべきものもありません。代わりにmemoryを支配するのがcacheであり、それは避けられないcostではなく、意図的な選択です。

fast runをもう一度見てください。最初のtokenは、残り47 tokensとは違う振る舞いをしていました。

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

promptはtokenあたり25.6 msで、生成tokenは1つあたり166 msでした。同じmodel、同じhardware、同じweightsなのに、tokenあたり6倍の差です。そして多くの人の予想とは逆方向です。promptが安い部分です。generationは、本当に異なる物理を持つ2つのphaseに分かれます。

両方のhalfは測定できます。Prefill、$P$ tokensに対する1 pass：

Decode、$C$のcacheに対して1 token：

2つ目の表は2回読んでください。contextが16 tokensから1,024へ、attendすべき履歴が64倍になっても、1 stepのcostは測定できるほど変わりませんでした。cacheに対するattentionは実際の作業ですが、1つのvectorを作るために5億個のweightsをmemory busに流し込む固定costに比べれば小さいのです。その固定costこそが、次のsectionのすべての理由です。

この2つのphaseが、serving systemが必ず報告する2つの数値の起源です。Time to first tokenは本質的にprefillであり、promptとともに増えます。だから長い会話は始まりが遅く感じます。Tokens per secondは$1/\text{decode step}$であり、ほぼ一定です。だからその後のreplyは均等に流れます。始まりが遅く、その後なめらかにstreamするchatは、renderingの工夫ではありません。この2つの表そのものです。

cacheは算術をmemoryと交換します。そして必要とするmemoryは小さくありません。context内の各tokenについて、各layerはkey-value headごとに1つのkey vectorと1つのvalue vectorを保持します。

2はkeysとvaluesの分です。それ以外はarchitectureです。この章全体で測っているmodel — 24 layers、14 query heads、2 key-value heads、head dimension 64 — では、fp16でtokenあたり$2 \times 24 \times 2 \times 64 \times 2 = 12{,}288$ bytesです。

この分野のformulaは2倍ずれる癖があるので、信じるのではなく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で一致し続けます。

最後の3 rowsはもう一度見る価値があります。32 usersで各2,048 tokens、64 usersで各1,024、128 usersで各512 — いずれの場合もcacheは768 MBです。3つとも65,536 tokensを保持しているからです。cacheはresidentなtoken総数だけに依存し、それがusersにどう分配されているかには依存しません。 この事実がbatching sectionの基礎です。

第9章ではmulti-query attentionとgrouped-query attentionを紹介し、その理由をこの章に先送りしました。理由はこのformula、とりわけその中の$H_{kv}$です。

標準的なmulti-head attentionでは、各query headが自分専用のkey headとvalue headを持ちます。ここでのmodelには14 query headsがあります。full multi-head attentionならcacheはtokenあたり$2 \times 24 \times 14 \times 64 \times 2 = 86{,}016$ bytes、つまり12 KBではなく84 KBになります。ちょうど7倍で、query headsとkey-value headsの比率そのものです。

Multi-query attention¹はこれを極限まで進め、すべてのquery headsが1つのkey-value headを共有します。Grouped-query attention²は勝ち残った妥協案です。少数のkey-value headsを持ち、それぞれをquery headsのgroupで共有します。MQAの品質低下は実在し、GQAではそうではなかったからです。どちらも算術を節約しません。このformulaを整数で割るために存在し、long contextsによってcacheがbinding constraintになった瞬間、業界全体に広がりました。

そしてそれはすぐにそうなります。32 layers、dimension 128のkey-value headsを8つ持つ7B-class modelでは、fp16のcacheはtokenあたり128 KBです。

そのmodel自身のweightsはfp16で13.0 GBです。この章の最後の表にある値です。したがって128,000-token contextでは、1 userのcacheがmodelより大きくなります。この算術を第16章ではお金に変換します。そして長い会話が単に遅いだけではない理由でもあります。requestが生きている間、その会話はmachineの固定された一部を占有します。

Decodeはmemory-boundです。weightsは1 tokenを作るためにbusを通って引きずられ、算術unitは遊んでいます。ならば同じstepにもっと仕事を入れます。複数requestを同時に走らせれば、1回読まれたweightsが全員に使われます。同じmodelで測ると、各requestが64-token cacheを持ち、1 tokenをdecodeする場合：

右2列を突き合わせて読んでください。そこがすべてです。request数を1から16に増やすと、throughputは6.0倍になりますが、個々のrequestの待ち時間は2.67倍になります。batchはserverを良くし、すべてのuserを悪くしました。

これはtuningで消せるbugではありません。それ自体がtrade-offであり、両側に名前があります。Latencyはreplyを待つ人が経験するものです。Throughputは請求書を割る相手です。どちらも改善するsettingはありません。

止まる場所にも注目してください。16から32ではthroughputの増加は9 %なのにlatencyはほぼ2倍になります。stepはmemory-boundではなくcompute-boundになっており、その膝を越えるとbatchは何も買いません。すべてのdeploymentにこの膝があります。場所はあなたの環境で測る必要がありますが、その存在は測るまでもありません。

単純なbatchingは、$B$ requestsを集め、一緒に走らせ、全員が終わったら返すというものです。しかし全員が同時に終わるわけではありません。replyが20 tokensのものもあれば500 tokensのものもあります。固定batchは最長のmemberが終わるまで走り、終了済みrequestもそれまでslotを占有し続け、paddingに寄与します。

output lengthに現実的な偏りがある64 requests — median 18 tokens、最長231、合計1,874 — を取り、8 slotsの測定済みper-step costで両policyをsimulateします。

throughputは1.6x改善します。平均latencyは10倍以上改善します。static batchingでは、4 stepsで終わったrequestでも、231-tokenの隣人が終わるまで誰にも届かないからです。

Continuous batching³が修正策で、聞こえるとおり単純です。batchはgroupではなくslotsの集合であり、空いたslotは次のstepですぐにqueued requestを受け入れます。schedulerはrequest単位ではなく1 token単位で動きます。本番のserving stackは今ではすべてこれを行います。

これには第2のhalfがあり、それがcacheです。出入りするslotsはcache memoryを断片化し、各slotに最大可能contextを予約すると、その予約の大半が無駄になります。PagedAttention⁴はoperating systemsから答えを借ります。cacheを固定サイズblocksに保存し、sequenceごとにblock tableを持つことで、sequenceのcacheは物理的には散らばっていても論理的には連続に見えます。これにより、prefixを共有する2つのsequencesがそのprefixを保持するblocksを共有することもできます。vLLMがその上に作られている理由であり、serving engineがtransformer付きのmemory allocatorである理由です。

請求書のもう半分はweightsそのものです。5億parametersは4 bytesずつなら1.98 GB、2 bytesなら0.99 GB、1 byteなら0.49 GBです。weightあたりのbitsを減らすと、modelはdisk上でもmemory上でも小さくなり、decodeはbandwidth-boundなので各stepも速くなります。動かすbytesが少ないからです。

最も単純な方式はsymmetric absolute-maximum quantizationで、3行に収まります。

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

最大weightが最大integerに写るようscaleを選び、割り、roundし、integersとscaleを保存します。復元は掛け戻すだけです。賢いことは何もしていません。そして動きます。動かなくなるところまでは。

modelの実weightsで測りました。168個すべてのprojection matrices、3億5,780万parameters、relative error $\lVert W - \hat{W}\rVert / \lVert W \rVert$：

4行目が崩壊です。worst matrixのrelative errorが0.99ということは、復元はoriginalをほとんど何も保持していないという意味です。そのmatrixは、だいたい正しい大きさのnoiseに置き換えられています。原因は、single 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 %)

6,000個に1つのweightが6 standard deviationsを超え、最大は24 outです。matrix全体で1 scaleを使うと、その1つのweightが430万個すべてのstep sizeを決めます。8 bitsなら256 stepsあり、典型的なweightも意味のあるlevelに落ちます。4 bitsでは16しかなく、外側のlevelはほとんど誰も持たない値のために予約され、普通のweights — つまりほぼ全員 — は2つか3つのlevelsに丸められます。

そのrow以降は、粒度を変えた同じ修理です。scaleの担当範囲を小さくします。output rowごとにするとerrorは3.4分の1になり、連続128 weightsごとのgroupにするとさらに下がります。costはbookkeepingです。128ごとのgroupに16-bit scaleを1つ置くと、weightあたり4 bitsではなく$4 + 16/128 = 4.125$ bitsになります。そして失った差の大半を取り戻します。

NF4は反対側から攻めます。⁵ levelsは等間隔である必要がありません。block内のweightsはおおよそnormal distributionなので、16 levelsをnormal distributionのquantilesとして選びます。weightsが実際にいるzero付近は密に、いないtailsは疎にします。同じ4 bits、同じblock scalingで、blockはより小さく — group-128の4.125に対してweightあたり4.25 bits — 測定errorは0.1323から0.0952へ、28 %下がります。その一部はより細かいblockによるもので、残りはmassのある場所にlevelsを置いたことによります。2つを分けるには3つ目のrowが必要です。

第2章のfloating-point boxは約束で終わりました。この章でweightsを8 bitsと4 bitsに量子化し、圧縮を拒む少数のoutlier featuresを見つける、という約束です。それがこれです。そして「数値を丸めるだけ」がactivationsに対してうまくいくはずがなかった理由も説明します。

上のweightsは行儀が悪いものでした。activationsは別格です。普通の84-token promptを取り、各layerのresidual streamをcaptureし、896 dimensionsそれぞれが到達する最大magnitudeを測ります。

dimension 62は1,579.6に達しますが、median dimensionは1.6を超えません。これは1 tokenや1 layerの偶然ではありません。同じdimensionがlayer 4にもあり、layer 20にもほぼ同じ値で残っています。これがoutlier features⁶であり、systematicです。inputではなくtrained modelの性質です。

layer 16での896個のper-dimension maximaのhistogramを見ると、その形は明白です。

     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

900 dimensionsが8未満にきれいに積まれ、3 octaves分は何もなく、最後の端に1 dimensionだけがいます。ではそのtensorをINT8に量子化し、何が起きるか数えます。

256 levels中14 levels。scaleは1,579.6で決まったため、各stepは幅12.44です。典型的なactivation — median magnitude 0.26、99th percentile 2.51 — には着地先がありません。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

1 level。 dimension全体、すべてのtokenが同じ数に量子化されました。8 bitsを割り当てたのに、使われたのはほぼ0 bitsです。そしてmodelはそれらのactivationsを読み、constantを渡されます。

この測定が、実際に使われるあらゆるtechniqueの根拠です。

quantizationについての記事のほとんどは前sectionで止まります。methodを説明し、compression ratioを引用し、品質は「ほぼ保たれる」と主張します。第4章は自分をだまさないための章だったので、実際に確かめます。

同じmodelのweightsを各schemeでin-placeに量子化し、3つを測ります。held-out English prose 2,048 tokensでのperplexity — ここではこのcourseのdraftです。そのためrepositoryでは固定のpublic-domain bookに差し替え、同じ形で異なる数値の表を出します — greedy decodingで既知の答えを持つ16の短い事実質問のbattery、そして同一contextを与えたときに量子化modelがfull-precision modelと一致するtokenの割合です。

この表で率直に述べるべきことが4つあります。

INT8は正しく行えば無料です。 per-row INT8はreferenceの23.08に対して22.96です。差は200分の1で、noiseであり「同一」と読むべきです。noiseがどちらを向くかは安定しません。repositoryのpublic-domain corpusでは同じ2 schemesが22.24対22.18になります。距離は半分で、向きは逆です。生成された144 tokensのうち142でfull-precision modelと一致します。fp32 referenceに対してmemoryは4分の1、実際にdeployするであろうfp16に対しては半分で、検出可能なcostはありません。INT8は雑にやってもほぼ無料です。matrixごとに1 scaleでもperplexityは0.5 points、battery answersは失いません。8 bitsは十分に寛容なので、granularityはほとんど問題になりません。だからこそ人はINT8からINT4へ一般化して痛い目を見ます。

tensorごとに1 scaleのINT4はmodelを破壊します。 Perplexity 3億6,500万。劣化ではなく壊滅です。その後はgranularityがすべてです。per-tensor 365,416,000、per-row 46.18、per-group-of-128 31.08、NF4 24.55。同じweightあたり4 bitsで、最悪と最良の間に1,500万倍の差があります。

Perplexityは粗いinstrumentで、batteryはさらに粗いものです。 NF4とgroup-128 INT4の間でperplexity差は6.5 points、battery差は1 questionです。そして第4章のconfidence intervalによれば、16問中1問の差は何も区別しません。intervalより鋭い実演もあります。model標準のrepetition penaltyをoffにして同じbatteryを走らせると、それがgreedy decodingの本来の意味ですが、この2 rowsは入れ替わります。16問中1問は小さなeffectではなく、effectなしです。第8章の警告も当てはまります。perplexityはtokenizerを共有するmodels間でしか比較できないため、他人の記事の数値を自分のものと比較することはできません。

agreement列は3つの中で最も鋭く、 ほぼ無料です。full-precision modelをgreedyに走らせ、各positionで同じprefixを与えたときにquantized modelなら何を選ぶかを尋ねます。16ではなく144の独立観測があり、ground truthを必要とせず、batteryが段差状に劣化する場所でもなめらかに劣化します。そしてそれは次sectionが必要とする量そのものでもあります。

これは第1章がこの章について約束したことの予定どおりの到着です。数学は4-bit modelが可能だと言い、engineeringはそれが使えるかどうかを決めます。

第12章でこれを予告し、請求書をここに残していました。

発想はprefill/decode分割からそのまま出てきます。$\gamma$ tokensの提案されたsequenceを検証するcostは、$\gamma$ positionsに対する1 forward passです。matrix-matrix productなので、1つに対するpassと比べてほとんど高くありません。つまり：

output distributionは変わりません。greedy decodingなら明らかです。targetが生成したはずのtokenだけをacceptするからです。samplingでは修正されたacceptance ruleが必要で、Leviathan et al.は得られるdistributionがtargetのものと完全に同じであることを証明しています。¹⁰ これはこの章で2つ目のexact optimizationです。

したがってすべてはacceptance rate $\alpha$にかかっています。これは測定可能です。上のagreement列であり、そこで計算した理由です。各quantized modelをfull-precision targetのdraftとして使い、144 generated positionsで測ると：

verification passごとにacceptされる期待tokens数は、draft length $\gamma$で

net speedupは、それをdraft自身のcostで割ります。draftのcostはtargetのtokenあたりcostに対する割合$c$です。

太字のentryを覚えてください。speculative decodingはgenerationを遅くすることがあります。 acceptanceが30 %で、draftがtargetの5分の1のcostなら、5回のforward passに支払って1.4 tokensしか保持できません。最後の列はもう1つの罠です。長いdraftはacceptanceが高いときにだけ役に立ちます。$\gamma$-token guessのtailにはほとんど到達しないからです。90 % acceptanceでは$\gamma = 8$は3.40xの価値があり、30 %では0.79xの価値しかありません。同じconfigurationが、あなたのtrafficで測った数値次第で勝ちにも負けにもなります。

Quantizationは同じfunctionをより少ないbitsで保存することでmodelを縮めます。Distillationは、小さなmodelを大きなmodelの模倣へtrainingすることで縮めます¹¹。deep learningよりほぼ10年古いideaです。¹²

微妙なのは、studentが何から学ぶかです。正解ではありません。それなら直接trainingできたはずです。teacherが加えるのは分布全体です。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もあり得て、large — adjectiveであり、文法的にはまったく別のcontinuation — もまだ生きていると教えます。これが元の議論です。これは7だが、1にもかなり似ている。その類似性は、hard labelが捨てる情報です。

だからdistillationはtemperatureを使います。softmaxの前にlogitsを$T$で割ると、distributionは平らになり、次点候補のrelative weightが上がります。このphraseでは、top tokenとthirdの比率が$T = 1$で2.24だったものから、$T = 2$で1.50に下がります。最初の値の平方根であり、logitsを2で割るとratioに起きることです。順序は同じで、near missesにlossのattentionがより多く向きます。studentのgradientはteacherの不確実性を運び、判定だけを運ぶのではありません。

この章のすべてはいま1つの和です。

ここで$T$は、同時request全体にまたがるresident tokensの総数です。適用してみます。7Bと70B rowsはdimension 128のkey-value headsを8つ仮定し、13B rowは40 headsのfull multi-head attentionを仮定します。その世代のmodelはそう作られていたからです。そしてそれは表に現れます。

8 GB

16 GB

24 GB

8 GB表の13B rowを見てください。weightsは収まります。8 GBのうち6.2 GBです。だから通常の言い方では、13B modelは「8 GB cardで動く」ことになります。contextは337 tokensです。これは会話ではなく、promptとしてもぎりぎりです。「収まるか」は間違った質問です。正しい質問は「どれだけのcontextで、何人のusersを同時に扱えるか」です。

16 GBのint8の2 rowsも見てください。7Bは65,378 tokens、13Bは3,136 tokensです。追加weights 5.6 GBから20倍の差が出ています。ここでの13Bはmulti-head attentionであり、そのcacheはtokenあたり800 KB、7Bの128 KBに対して大きいからです。似たsizeの2 modelsの一方がlong contextで使い物にならない。その理由はどのmodel cardのheadlineにも出てきません。

13章前、これは2つのweightsとbiasを持つperceptronでした。いまや、設計され、trainingされ、alignedされ、難しい質問にcomputeを使うよう教えられ、tokenあたり測定済みcostでservedされるtransformerです。中に未開封のboxはもう残っていません。

ここで終わります。そして意図的に終わります。

第14章は、modelが別の場所にあるところから始まります。あなたのprocessの中でも、memoryの中でも、printできるvariableの中でもありません。あなたが管理しないmachine上、API keyとportと請求書の向こう側です。ここで測ったすべてはまだ起きています。first tokenの前にはprefillが走り、cacheは会話とともに増え、あなたが入っているbatchは依然として誰か別のものであり、あなたのlatencyを決めています。しかしこれからは、それをServer-Sent Eventsのstream、finish_reason、そしてRetry-After header付きHTTP 429を通して観測します。視点が変わると問いも変わります。このgradientはどう計算されるのかではなく、なぜ請求額が3倍になったのかになります。言葉も変わります。第14章はそのruleを宣言ではなく説明します。ここまではcodeがweights、gradients、logits、tokenizer bytesを持っていました。そこから先ではconnection、retry、cancellation、accumulated stateを持ちます。この境界を越えても、背後にある13章は捨てられません。それらはportの向こう側で動いているものの説明です。

2つの省略は意図的です。FlashAttention (Dao et al., arXiv:2205.14135) は別のattentionではありません。同じfunctionを、operationをtile化して$n \times n$ score matrixをmemoryに書き出さないように計算します。だから実践では、この章の2つ目の表にある67 MBは算術が示すより小さくなります。そしてkernels自体は委譲しています。Stanford CS336のlecture 10はinference systemsを、ここでは試みない深さで扱っています。CPU側については、llama.cpp repositoryとGGUF specificationがprimary sourcesです。