CUDA Graphsとは？LLM推論のkernel起動 overheadを減らす仕組み

LLMのDecodeでは、一tokenを生成するたびに多数の小さなGPU kernelが起動する。GPU計算そのものが短くなるほど、CPUとdriverがkernelを準備してsubmitする時間の割合が大きくなる。

CUDA Graphsは、繰り返すGPU処理と依存関係を一度graphとして定義し、事前準備した実行形を何度もlaunchする仕組みである。kernel launchのたびに同じ準備を繰り返す量を減らし、短いiterationのCPU overheadとjitterを抑えられる。

ただし、requestごとにshapeや処理経路が変わるLLM servingでは、何でも一つのgraphへ固定できるわけではない。capture可能な範囲、shape bucket、固定memory address、fallback経路を設計して初めてproductionで使える。

streamへ一つずつsubmitするcost

通常のCUDA streamでは、applicationがkernelやmemory copyを順にenqueueする。各操作についてhost driverがlaunch準備を行うため、短いkernelを大量に繰り返すworkloadではhost側のoverheadが目立つ。

NVIDIAのCUDA Programming Guideは、CUDA Graphをoperationとdependencyの集合として定義し、definition、instantiation、executionを分けている。

definitionでkernel、copy、依存関係をgraphとして作る
instantiationでgraphを検証し、実行準備を行う
executable graphをstreamへ何度もlaunchする

事前準備を一度へ寄せることで、繰り返しlaunchのcostを減らす。graph全体をCUDAへ渡すため、operationを一つずつsubmitする場合には難しい最適化も可能になる。

LLM Decodeと相性がよい理由

Decodeは、一度に生成するtokenが少なく、同じlayer列をtokenごとに繰り返す。batch shapeとmemory位置が安定していれば、同じkernel sequenceを何度も実行する形になる。

Prefillは入力token数が大きく、matrix演算も長い。GPU計算時間が十分長ければlaunch overheadの比率は小さくなる。対して小batchのDecodeでは一つ一つのkernelが短く、CPU submitがGPUの次の仕事に間に合わない状態が起きやすい。

CUDA Graphsが効くかは「LLMだから」ではなく、同じ短いGPU workflowをどれだけ繰り返し、host launchがbottleneckになっているかで決まる。

graphはoperationと依存関係を持つ

graphのnodeにはkernel、memory copy、memset、event待機、external semaphoreなどがある。edgeはnode間のdependencyを表す。upstream nodeが終わった後にdownstream nodeを実行できる。

作成方法は二つある。

Graph APIでnodeとdependencyを明示的に追加する
既存streamへ発行した処理をstream captureで記録する

frameworkやLLM engineでは、既存のCUDA処理をcaptureする方式が使われることが多い。ただしcapture中に対応しないoperation、CPU依存の分岐、synchronizationがあれば失敗する。

capture時のCPU codeはreplayされない

PyTorch向けのNVIDIA資料は、graph replayではcaptureされたGPU kernelだけが実行され、capture範囲内にあったPython logic、CPU計算、I/Oは再実行されないと説明している。

たとえばcapture中のPython counterでkernel parameterを決めても、replay時にcounterを更新しただけではGPU処理へ反映されない。capture時の値が固定され、古い条件のまま実行される危険がある。

sampling、stop判定、request終了、tool callへの分岐など、tokenごとにCPU判断が必要な処理を無理にgraphへ含めない。GPU上で固定して繰り返せる部分と、iterationごとに判断する部分を分ける。

memory addressの安定が必要である

graph内のkernelは、capture時に使ったtensorのmemory addressを参照する。replayごとに新しいtensorをallocateしてpointerが変わると、そのままでは期待したdataを読めない。

そのため、input・output用のstatic bufferを確保し、replay前に新しいdataを同じbufferへcopyする構成が使われる。frameworkのgraph用memory poolは、captureされたallocationのaddressを安定させる。

LLM servingではweightだけでなく、activation、KV Cache block table、sequence metadata、sampling parameterの渡し方を確認する。memory addressが同じでも、長さや有効要素数が誤っていれば古いdataを計算へ混ぜる。

dynamic shapeはgraphを増やす

request数、prompt長、Decode中のactive sequence数は変化する。一つのexecutable graphは、capture時のworkflowとshapeへ強く結び付く。

実務では、代表的なbatch sizeやtoken数ごとにgraphをcaptureするshape bucketを使う。たとえば1、2、4、8、16 request向けのgraphを用意し、実batchを近いbucketへpaddingする。

bucketを増やせばpadding wasteは減るが、capture時間とgraph用memoryが増える。bucketを減らせば運用は単純になるが、小さなrequestを大きなshapeで実行してcomputeを無駄にする。

監視では実batch sizeだけでなく、選択されたgraph bucket、padding率、fallback率を見る。

full graphとpiecewise capture

LLM iteration全体を一つのgraphへcaptureできればlaunch数を大きく減らせる。しかし、dynamic operationや未対応kernelが一つあるだけでfull captureが難しくなる。

そこで、capture可能なregionだけをgraph化し、その前後を通常実行するpiecewiseな方法がある。overhead削減幅は小さくなるが、dynamicなmodelや変更の多いserving codeへ適用しやすい。

full graphかpiecewiseかを機能名だけで比較しない。実際のlaunch数、CPU time、fallback箇所、graph境界でのsynchronizationをprofileする。

warmupとcaptureは最初のrequestへ載せない

graphのinstantiationはvalidationと実行準備を行う。最初のmodel load直後にはkernel compilation、memory allocation、autotuningなども起きうる。

いきなり本番requestをcapture対象へ流すと、初回だけ大きなTTFTが出る。deployment readinessの前に、予定するshape bucketをwarmupし、captureとinstantiateを終える。

一方、すべての想定shapeを起動時にcaptureすると、startup時間とmemoryを使う。頻度の高いbucketをeagerに準備し、低頻度bucketを通常実行または遅延captureする設計もある。

Scale to Zeroとの組み合わせでは、GPU node起動、model load、compile、graph captureの合計がcold startになる。Podが起動した時点をreadyとしない。

graph replayが常に速いわけではない

CUDA Graphsが減らすのは主にhost launch overheadである。大きなmatrix計算、memory bandwidth不足、collective通信、KV Cache missがbottleneckなら、graph化だけでは改善しない。

さらに、paddingで余分なtokenを計算すれば、launch削減よりGPU計算増加が大きくなる場合がある。graphあり・なしを同じ実traffic分布で比較する。

見るべき結果はmicrobenchmarkのiteration時間だけではない。

TTFTとITLのp50、p95、p99
token throughputとGPU当たりrequest数
CPU core利用率とlaunch待ち時間
GPU idle gap
graph hit・miss・fallback率
paddingされたtoken数
capture・instantiate時間と追加memory

model更新とgraph invalidation

model version、quantization、kernel、attention backend、parallelism、max sequence、KV Cache layoutが変わると、以前のgraphを再利用できない場合がある。

parameterだけが変わりgraph topologyが同じ場合、CUDAにはexecutable graphをupdateする仕組みがある。しかしnode typeや依存構造が大きく変われば再instantiationが必要になる。

LLM engineが内部でgraph cacheを管理していても、deploymentをまたいで安全に再利用できると仮定しない。model artifact、engine version、driver、CUDA、GPU architecture、shape設定をgraph cache keyへ含める。

更新失敗時は古いgraphを誤ってlaunchせず、再captureまたは通常実行へfallbackする。

NCCL operationもcapture対象になりうる

NVIDIA NCCLはCUDA Graphsと組み合わせ、compute kernelとcollective operationを同じgraphへcaptureする方法を提供している。Tensor Parallelismなどで毎iteration同じcollectiveを行う場合、通信launchのoverheadもまとめられる可能性がある。

ただし、全rankが同じ順序とshapeでcollectiveへ入る必要は変わらない。一部rankだけ別graphやfallback pathへ入れば、collective mismatchやhangを起こしうる。

rankごとのgraph ID、bucket、iterationをlogへ残す。distributed graph captureとreplayは、一台の成功だけで確認せず全rankで試験する。

Continuous Batchingとの関係

Continuous Batchingでは完了したsequenceを外し、新しいrequestをiterationごとに入れる。active sequence数が変わるため、CUDA Graphsが好む固定shapeと緊張関係がある。

解決策はbatch size bucket、padding、metadataのstatic buffer化、未対応shapeのfallbackである。graphのためにschedulerを硬直させ、終了requestを長く残せばlatencyとcapacityを損なう。

graph hit率だけを目標にしない。Continuous Batchingの入れ替え効率、KV Cache utilization、fairnessと合わせて最適化する。

Speculative Decodingとの関係

Speculative Decodingでは一iterationに検証するtoken数が変わり、draftのacceptance結果によって後続処理も変化する。固定した1-token Decodeよりgraph variantが増えやすい。

検証token数やbatch sizeごとにbucketを作ると、graph数とwarmup costが増える。CUDA Graphsを有効にしたことでSpeculative Decodingのacceptance benefitが消えていないか、組み合わせごとに測る。

監視するmetric

graph capture、instantiate、uploadの時間
graphごとのmemory使用量
graph launch数と通常kernel launch数
shape bucket別のhit、miss、fallback
padding token数と実token数
graph replay前後のCPU利用率
GPU idle gap、kernel duration、memory bandwidth
TTFT、ITL、throughputの分布
capture failure、update failure、re-instantiation回数
model・engine・driver version別のgraph cache
rank別のgraph ID、bucket、NCCL timeout

回帰試験

graphあり・なしで同じoutputを比較する
batch sizeとsequence lengthの境界値を通す
bucket外shapeが安全にfallbackするか確認する
request終了と新規追加を繰り返す
sampling parameterとstop条件がrequest間で混ざらないか確認する
static bufferに前requestのdataが残っても結果へ混ざらないか確認する
model、quantization、kernel変更後に古いgraphを使わないか確認する
startup直後のwarmupとreadinessを確認する
multi-rankで一つのrankだけfallbackした場合を検出する
eager executionへ切り戻してserviceを継続できるか確認する

性能試験では平均だけでなく、短いrequestのp99を見る。graph captureや再instantiationが本番traffic中に発生する場合は、そのspikeも含める。

繰り返しが安定している部分だけを固定する

CUDA Graphsは、GPU operationとdependencyを一度準備し、同じworkflowを低overheadでreplayする仕組みである。小batch Decodeのように短いkernelを繰り返す処理では、CPU launch overheadとjitterを減らせる。

一方、dynamic shape、CPU分岐、変動するmemory addressはgraph化を難しくする。shape bucket、static buffer、warmup、fallbackを設計し、graph hit率とpadding wasteを同時に見る必要がある。

速さはgraphを有効にした事実では決まらない。実trafficでTTFT、ITL、throughput、CPU time、GPU idle gapを測り、繰り返しが安定した部分だけを固定することが重要である。