PagedAttentionとは？LLMのKV Cacheをblock単位で管理する

LLM推論では、model weightをGPUへ載せた後の空きmemoryをKV Cacheとして使う。同じmodelでも、同時に処理できるrequest数はpromptと回答の長さで変わる。各requestの最大長を見込んで大きな連続領域を先に確保すると、使われないmemoryが増え、短いrequestを追加できなくなる。

PagedAttentionは、requestのKV Cacheを固定長のblockへ分け、論理上の並びとGPU上の物理的な配置をblock tableで対応させる。OSのvirtual memoryとpagingに着想を得て、連続していない物理blockを一つのcontextとしてAttentionから読めるようにした仕組みである。

KV Cacheは生成tokenとともに増える

Transformerは次tokenを生成するとき、過去tokenのKeyとValueをAttentionで参照する。毎回すべてを再計算しないため、そのKeyとValueをKV Cacheへ保存する。

requestが1 token進むたびにCacheも増える。問題は、生成前には最終的なoutput長が分からないことである。max_tokens=2,000を指定しても、50 tokenで終了するかもしれない。反対に長い回答が上限近くまで続く場合もある。

requestごとに最大長分の連続GPU memoryを予約すればOOMを避けやすいが、実際には使わない領域を抱える。必要量だけを都度拡張しようとしても、十分な総空き容量があるのに連続領域が見つからないexternal fragmentationが起こりうる。

tokenを固定長blockへ分ける

PagedAttentionでは、KV Cacheを複数token分の固定長blockで管理する。説明を簡単にするため、一つのblockへ4 token分を保存するとする。

12 tokenのrequestには論理block 0、1、2が必要である。GPU上では物理block 7、2、11へ分散していてもよい。

論理block	token位置	物理block
0	0〜3	7
1	4〜7	2
2	8〜11	11

block tableがこの対応を持つ。Attention kernelは論理順にtableを参照し、分散した物理blockからKeyとValueを読む。物理memory上で連続した大領域をrequestごとに確保する必要がない。

生成が4 token増えれば、新しい空き物理blockを一つ割り当てる。requestが終了すれば、そのrequestが参照していたblockをpoolへ戻し、別requestが再利用できる。

memory wasteは最後のblockへ限定される

block単位で必要な分だけ割り当てれば、requestの将来最大長を先に予約しなくてよい。物理blockの間に空きができても、ほかのrequestへ割り当てられるためexternal fragmentationを抑えられる。

一方、requestのtoken数がblock sizeの倍数でなければ、最後のblockには未使用slotが残る。これがinternal fragmentationである。たとえばblock size 16で17 tokenなら、2個目のblockは1 tokenだけ使い、15 slotが空く。

blockを小さくすれば最後の無駄を減らせるが、block tableのentry数と管理回数が増える。大きくすればmetadataは減るが、短いrequestの最後で無駄が増える。最適なblock sizeはmodel、kernel、hardware、request長の分布で決まる。

PagedAttentionはAttention kernelとmemory managerの組み合わせである

単にKV Cacheを小分けにするだけではない。通常のAttention kernelが連続tensorを前提としていれば、分散blockをそのまま読めない。PagedAttentionはblock tableに従ってKeyとValueを読むkernelと、blockの割り当て・解放・参照数を管理する仕組みを組み合わせる。

vLLMのPagedAttention実装は、KV Cacheの物理blockとslot mappingを使ってKeyとValueを書き込む。schedulerはrequestが次に必要とするblockを確保し、model runnerがその対応をkernelへ渡す。

そのため、PagedAttentionは一般的な「memoryをpageへ分けるアイデア」だけでなく、serving engine全体のscheduleと連動する。

OSのpagingと同じではない

着想はvirtual memoryだが、仕組みをそのまま同一視しない方がよい。

OSではvirtual addressをpage tableでphysical memoryへ写し、必要に応じてdiskとのpage in・page outも行う。PagedAttentionの中心は、token順の論理KV blockをGPU上の物理blockへ対応させ、Attention kernelから参照することである。

CPUやSSDへKV Cacheをoffloadする実装もあるが、それはpagingという名前から自動的に得られる性質ではない。転送経路、connector、帯域、復帰latencyを別に設計する必要がある。

Copy-on-Writeで共通prefixを共有できる

同じpromptから複数の候補を生成するparallel samplingやbeam searchでは、分岐前のprefixが共通である。候補ごとに同じKV Cacheを複製するとmemoryを浪費する。

blockへ参照数を持たせ、共通prefixを複数sequenceから参照すれば、分岐前のblockを共有できる。どれかのsequenceが共有blockを書き換える必要が生じた時点で、新しい物理blockへcopyする。これはCopy-on-Writeの考え方である。

PagedAttention論文では、block mappingとCopy-on-Writeによってparallel samplingのsequence間でKV Cacheを共有する設計を示している。

Prefix Cachingとの違い

PagedAttentionとPrefix Cachingは関係するが、同じ機能ではない。

PagedAttentionは、実行中requestのKV Cacheを物理memoryへどう配置するかを扱う
Prefix Cachingは、完了済みまたは別requestで計算した共通prefixのKV blockを、後続requestが再利用する

Prefix Cachingでは、同じtoken列とmodel条件から作られたblockを特定する必要がある。vLLMのAutomatic Prefix Caching設計は、親blockのhash、block内token、LoRA ID、multimodal入力hashなどを使ってblockを識別する。

PagedAttentionが効率的なblock allocationの土台となり、その上でhash-based lookupやevictionを加えたものがPrefix Cachingである。PagedAttentionを使うだけで、別requestのprefix計算が自動的に省略されるわけではない。

blockが足りなくなったとき

GPUのKV Cache poolは有限である。実行中requestが増え、すべての物理blockが使用中になれば、新しいtoken用blockを割り当てられない。

serving engineは状況に応じて、requestの受付を待たせる、実行中requestをpreemptする、KV CacheをCPUへ退避する、後でPrefillを再計算するなどの判断を行う。どの方法も無料ではない。

queueで待たせるとTTFTが伸びる
CPU swapは転送latencyとCPU memoryを使う
recomputeはGPU計算を繰り返す
requestの強制終了は利用者へerrorを返す

PagedAttentionはfragmentationを減らすが、物理容量そのものを増やさない。長いcontextと多い同時実行が重なればcapacity上限へ達する。

hybrid modelではblock設計が複雑になる

すべてのlayerが同じfull attentionを使うmodelなら、layerごとのKV Cache形状を揃えやすい。近年はfull attention、sliding window attention、Mamba系layerなどを組み合わせるmodelもある。layerによって必要なCache量と保持範囲が異なる。

vLLMのHybrid KV Cache Manager設計は、同じattention typeとpage sizeを持つlayerをgroup化し、groupごとにblockを割り当てる。sliding windowで古いtokenを解放する場合と、full attentionですべてを保持する場合ではcache hitの探索とevictionも異なる。

したがって、modelを変更したときに「block sizeが同じだからcapacityも同じ」とは限らない。attention architectureを含めて測り直す。

監視するmetric

GPU memory使用率だけでは、block poolが健全か分からない。model weightが一定なら、KV Cache領域の内部で何が起きているかを分けて見る。

GPU・CPUのKV block総数、使用中、free、cached
KV Cache使用率とtoken capacity
requestごとの割当block数
Prefix Cache hit token数とhit率
block eviction数と再利用までの時間
preemption、swap、recomputeの回数と時間
OOMとallocation failure
running・waiting request数
prompt長、output長、同時実行数の分布
TTFT、ITL、E2E latency

free blockが少ない状態が続けば、burst時にpreemptionが起きやすい。平均cache使用率ではなく最小free blockとtail latencyの関係を見る。

負荷試験

同じ長さのrequestだけではfragmentationと共有の効果を評価しにくい。次を混ぜる。

短文とcontext上限付近の長文
早くEOSへ到達する生成とmax tokenまで続く生成
同じprefixを持つrequestと完全に異なるrequest
parallel samplingで共通prefixから複数候補を生成
Prefix Cacheをcold・warmへ切り替える
block poolを圧迫し、preemptionやrecomputeを発生させる
requestを途中cancelし、blockが解放されるか確認する
model version、LoRA、multimodal入力を変え、誤共有しないか確認する

throughputだけでなく、peak同時実行数、waste、preemption、TTFTとITLを比較する。

PagedAttentionは限られたKV Cacheを細かく使う基盤である

PagedAttentionは、増え方の読めないKV Cacheを固定長blockへ分け、論理blockと物理blockを対応させる。連続領域の事前予約とexternal fragmentationを抑え、request終了時の解放、共通prefixのCopy-on-Write共有を行いやすくする。

ただし、GPU memoryの物理上限をなくす技術ではない。block size、Prefix Caching、preemption、offload、schedulerを含めて運用し、長さの異なる本番requestでcapacityとlatencyを測る必要がある。