Roofline Modelとは？LLM推論が計算とmemoryのどちらで詰まるか見分ける

GPUのFLOPSが二倍の製品へ変えたのに、LLMのtoken生成速度がほとんど上がらないことがある。原因は計算器ではなく、model weightやKV Cacheをmemoryから運ぶ速度が上限になっているからかもしれない。

Roofline Modelは、処理がcompute-boundなのかmemory-boundなのかを、演算量、data移動量、hardwareのpeak compute、memory bandwidthから見る性能modelである。細かな実装をすべて再現するものではないが、どの最適化へ時間を使うべきか絞り込める。

二つの屋根で性能上限を見る

Roofline Modelでは、ある処理が達成できる性能を概ね次の小さい方と考える。

P_{attainable} = \min(P_{peak},\ B_{memory} \times I)

$P_{peak}$ はhardwareのpeak演算性能、 $B_{memory}$ はmemory bandwidth、 $I$ はarithmetic intensityである。

arithmetic intensityは、memoryから運ぶ1 byte当たりに何回の演算を行うかを表す。

I = \frac{\text{FLOP}}{\text{byte transferred}}

演算数が少なくdata移動が多い処理は左側に位置し、memory bandwidthの斜めの屋根に制約される。演算数が十分多い処理は右側へ移り、peak computeの水平な屋根に制約される。

ridge pointが境界になる

memory-boundとcompute-boundの目安になるarithmetic intensityは、hardwareのpeak computeをmemory bandwidthで割って求める。

I_{ridge} = \frac{P_{peak}}{B_{memory}}

処理の $I$ がridge pointより小さければmemory-bound、大きければcompute-boundになる可能性が高い。

同じGPUでもdtypeによって $P_{peak}$ は変わる。FP32、BF16、FP8、INT8でridge pointは同じではない。Tensor Coreを使えるshapeか、sparse性能を前提にしているかでも公称値が変わる。

spec sheetの最大値を混ぜず、実際に使うprecisionとexecution pathで計算する。

matrix multiplicationの算術強度

$M \times K$ のmatrixと $K \times N$ のmatrixを掛けると、おおよそ $2MKN$ FLOPを行う。inputとoutputを一度ずつmemoryから読み書きすると単純化した場合、arithmetic intensityはshapeが大きいほど高くなる。

NVIDIAのMatrix Multiplication Background Guideは、matrix-vector productのように $M$ または $N$ が1に近い処理はarithmetic intensityが低く、memory-limitedになりやすいと説明している。大きなmatrix同士の積は、同じdataを多くの演算で再利用できるためcompute-limitedへ近づく。

ただし、この見積もりはcache、pointer演算、kernel実装、alignmentなどを省いた一次近似である。

Prefillは計算をまとめやすい

Prefillではprompt内の多数tokenを並列に処理する。linear layerは大きなmatrix multiplicationになり、同じweightを多くのtokenへ使える。

weightをmemoryから一度読んで複数tokenの演算へ再利用できれば、1 byte当たりの演算数が増える。promptやbatchが大きいPrefillはDecodeよりcompute-boundへ近づきやすい。

そのためPrefillでは、Tensor Core利用率、matrix shape、tiling、kernel効率が重要になる。GPUの演算性能を上げる効果も出やすい。

ただし、長文attentionではQ、K、Vや中間dataの移動が増える。すべてのPrefill operatorがcompute-boundとは限らない。operator単位で見る必要がある。

Decodeはweightを繰り返し読む

autoregressive Decodeは、一stepで通常一tokenずつ処理する。batchが小さいとlinear layerはmatrix-vectorに近くなり、巨大なweightを読み込んでも行う演算が少ない。

次tokenでも同じweightを再び使うが、model全体をon-chip cacheへ保持できるわけではない。各stepでHBMからweightを読むcostが支配的になりやすい。

この状態でpeak FLOPSだけ高いGPUへ替えても、memory bandwidthが同程度ならtoken/sは伸びにくい。DecodeでGPUのcompute utilizationが低いことは、必ずしもengineが壊れている証拠ではない。memoryの屋根へ達している可能性がある。

batchを増やすとweightを再利用できる

Decode batchを増やすと、一度読んだweightを複数sequenceのtoken計算へ使える。weight byte当たりのFLOPが増え、arithmetic intensityが上がる。

これがContinuous Batchingでthroughputを上げられる理由の一つである。少数requestではmemory-boundでも、多数requestをまとめるとcompute-boundへ近づく。

一方、batchを増やせば待ちqueue、KV Cache、ITLが増える。Roofline上で右へ移動することと、利用者のlatency SLOを守ることは別である。

最大throughputを得るbatchではなく、TTFTとITLを守れるoperating pointを探す。

KV Cacheもmemory trafficになる

Decodeのattentionは、過去tokenのKey・ValueをKV Cacheから読む。contextが長くなるほど、一token生成で読むKV量が増える。

短いcontextではweight readが主なmemory trafficでも、長いcontextではKV Cache readが大きな割合を占める。Grouped-Query AttentionやMulti-Query Attention、KV Cache量子化はKV byteを減らし、memory-boundなattentionを改善する可能性がある。

ただし、dequantization kernelやpage table accessが増える。byte数だけでなく実kernelのmemory transactionとlatencyを測る。

weight量子化はRoofline上の位置を変える

BF16 weightをINT8やINT4へ縮めれば、同じparameterを読むbyte数が減る。Decodeがweight bandwidthに制約されている場合、arithmetic intensityは上がり、token/s改善が期待できる。

しかし、量子化の効果は単純にbit幅の比にはならない。

scaleとzero pointも読む
unpackとdequantizeの演算が増える
対応kernelでなければtemporary bufferを作る
一部layerが高precisionのまま残る
small shapeでTensor Coreを使い切れない

memory traffic削減と追加computeの両方をRooflineで考え、品質評価も別に行う。

Tensor Parallelismはcomputeと通信を変える

Tensor Parallelismでweightを複数GPUへ分けると、一GPUが読むweightと行う演算は減る。ただしAll-ReduceやAll-Gatherが加わり、network・NVLink bandwidthという新しい屋根が現れる。

単一GPUのRooflineだけでは、collective待ちやPCIe・NVLinkの上限を表せない。分散推論では少なくとも次を分ける。

HBM bandwidthによるmemory roof
Tensor Coreなどによるcompute roof
NVLink、PCIe、networkによるcommunication roof
CPU launchやsynchronizationによるlatency

GPUを増やして一GPU当たりのGEMMが小さくなると、arithmetic intensityやkernel効率が下がる場合もある。

Kernel Fusionはdata移動を減らす

前稿のKernel Fusionは、operator間の中間tensorをHBMへ書かず、registerやshared memoryで次の演算へ渡す。これはRooflineの分母であるbyte trafficを減らす最適化と考えられる。

FlashAttentionも、大きなattention matrixをHBMへmaterializeせず、tilingでon-chip memoryを活用する。演算式だけでなくmemory hierarchyに合わせてdata movementを減らす。

memory-boundな処理へcomputeだけを足すより、fusion、tiling、layout、quantizationでbyteを減らす方が効く場合がある。

Speculative Decodingは一回のweight readで仕事を増やす

通常Decodeではtarget modelが一stepに一tokenを確定する。Speculative Decodingではdraft modelが複数候補を出し、target modelがまとめて検証する。

候補をまとめて検証できれば、target modelの一回のweight readに対する演算量が増え、memory-boundな小batch Decodeをよりcompute-intensiveな処理へ変えられる可能性がある。

ただし、実際の効果はacceptance rate、draft cost、追加KV Cache、verification shapeで決まる。Rooflineは上限を考える道具であり、採択率を含むend-to-end latencyは別に測る。

Rooflineは実測性能を保証しない

基本Rooflineはpeak computeとpeak bandwidthの二つだけを見る。実際には次の要因で屋根より低くなる。

matrix dimensionとTensor Core alignment
occupancy、register pressure、shared memory
kernel launch overheadとCPU bottleneck
cache hit率とmemory access pattern
branch、mask、padding、load imbalance
collective通信とrank synchronization
thermal・power throttling
frameworkのgraph breakやfallback

rooflineより大きく下にいるからといって、すぐにmemoryかcomputeのどちらかだけを増やさない。profilerでkernel、memory、communication、CPU timelineを確認する。

公称値ではなく実効roofを測る

vendorのpeak FLOPSとmemory bandwidthは理論上限である。productionで利用できるroofは、dtype、clock、power limit、GPU partition、同居workload、kernelで変わる。

memory bandwidth benchmark、GEMM benchmark、collective benchmarkを同じnode、driver、power設定で測り、実効値を使う。MIGやMPSで共有している場合は、そのpartitionと競合条件で再測定する。

GPU型番ごとの一般論より、自分のdeploymentのmeasured roofを持つ方が役立つ。

監視するmetric

operator・kernelごとのFLOPとmemory byte
arithmetic intensityとridge point
achieved FLOPSとpeak比
HBM bandwidthとpeak比
Tensor Core utilization、SM utilization
batch、prompt長、context長、active sequence数
weight、KV Cache、activationのread・write量
Prefill TTFT、Decode ITL、token throughput
NCCL collective時間とlink bandwidth
kernel launch gap、CPU利用率
dtype、quantization、GPU profile、clock

平均値だけでなく、workload class別に分ける。短文chatと長文RAGでは同じmodelでもRoofline上の位置が違う。

性能試験

PrefillとDecodeを別々に測る
batchを1から増やし、memory-boundからの変化を見る
prompt長とcontext長を独立に変える
BF16、FP8、INT8・INT4でbyteと品質を比較する
KV Cache precisionとattention backendを変える
fusionとCUDA Graphsを個別に有効化する
Tensor Parallel sizeを変え、communicationを含める
cold startとsteady stateを分ける
power・temperature制限時の実効roofを測る
SLOを守る最大throughputを記録する

bottleneckに合う最適化を選ぶ

Roofline Modelは、演算性能とmemory bandwidthから、処理が到達できる性能上限を考える。arithmetic intensityが低ければmemory-bound、高ければcompute-boundになりやすい。

LLM推論では、tokenをまとめて計算するPrefillはcompute-boundへ近づき、小batchでweightを繰り返し読むDecodeはmemory-boundになりやすい。batch、量子化、KV Cache、fusion、Speculative Decodingは、演算数とdata移動の比を変える。

重要なのはGPUの公称FLOPSだけで選ばないことである。どのphase、operator、traffic classがどの屋根へ当たっているか測り、そのbottleneckに効く最適化を選ぶ必要がある。