Kernel Fusionとは？LLM推論でGPU memory往復と起動回数を減らす

GPU使用率が高いのにLLM推論が思ったほど速くならないとき、計算量だけでなくdataの移動を見る必要がある。小さな演算を別々のGPU kernelとして実行すると、途中結果を何度もHBMへ書き戻し、次のkernelがまた読み込む。kernel起動も演算ごとに必要になる。

Kernel Fusionは、連続する複数の演算を一つ、または少数のGPU kernelへまとめる最適化である。途中結果をregisterやon-chip memoryに保ち、HBMとの往復とlaunch回数を減らせる。

ただし、演算をまとめれば常に速くなるわけではない。register pressure、parallelism、dynamic shape、数値精度、compile時間まで含めて判断する必要がある。

GPUは計算だけをしているわけではない

LLMの一つのlayerにはmatrix multiplication、bias、activation、normalization、residual addなどが並ぶ。eager executionで各operatorを別kernelとして起動すると、概念的には次の流れになる。

kernel AがinputをHBMから読む
中間結果をHBMへ書く
kernel Bが中間結果をHBMから読む
次の中間結果をHBMへ書く
kernel Cが再び読む

演算量が小さいelement-wise operationでは、算術よりmemory accessの時間が支配的になりやすい。fusionによってA、B、Cを一つへまとめられれば、中間結果をon-chipに保ち、最終結果だけをHBMへ書ける場合がある。

fusionが減らす二つのcost

一つ目はmemory trafficである。中間tensorをmaterializeせず、registerやshared memoryで次の演算へ渡せれば、HBM read・writeが減る。

二つ目はkernel launch overheadである。三つのkernelを一つへまとめれば、CPUとdriverからのlaunchも一回になる。前稿のCUDA Graphsもlaunch overheadを減らすが、方法が違う。

Kernel Fusionは複数operatorを一つのkernel実装へまとめる
CUDA Graphsは複数kernelを含む同じworkflowを事前準備してreplayする

fusion後のkernel sequence全体をCUDA Graphへ入れることもできる。両者は排他的ではない。

典型的なfusion

LLM推論で見かける代表例には次がある。

linearの出力へbiasとactivationを続ける
residual addとlayer normalizationをまとめる
quantized weightのdequantizationとmatrix multiplicationをまとめる
rotary position embeddingの変換をまとめる
attention内のscale、mask、softmax、value集約をまとめる
optimizer更新をまとめる。ただしこれは主に学習時である

どの組み合わせが使えるかはGPU世代、dtype、shape、framework、backendで変わる。名前が同じmodelでもkernel構成が同じとは限らない。

FlashAttentionはIOを中心に設計する

FlashAttentionは単にattentionの式を近似する方法ではない。原論文は、GPUのHBMとon-chip SRAMの間のread・writeを減らすIO-awareなexact attention algorithmとして提案している。

標準的な実装で大きなattention score matrixをHBMへmaterializeすると、sequence lengthが長いほどmemory trafficと容量が増える。FlashAttentionはtilingを使い、Q、K、Vのblockをon-chip memoryへ読み込み、softmaxを分割して計算しながら結果を蓄積する。

attentionの数学的意味を変えずにdata movementを減らす点が重要である。ただしfloating-point演算の順序は変わりうるため、bit単位で同一とは限らない。許容誤差と最終品質で確認する。

torch.compileはoperator列をgraphとして見る

PyTorchのtorch.compileは、Python programから取得したgraphをbackendへ渡し、TorchInductorなどが最適化したkernelを生成する。operatorを一つずつしか見ないeager実行より、operator境界をまたいだfusionを行いやすい。

公式APIには、pointwise operationをtemplateへ融合するepilogue_fusion、matrix shapeをTensor Coreに合いやすくするshape_padding、CUDA Graphsを使うoptionなどがある。つまりcompile、fusion、autotune、CUDA Graphsは別概念だが、一つのbackend設定で組み合わされる場合がある。

性能差の原因を調べるときは「compileを有効にした」だけで終わらず、生成graph、fusion後のkernel数、graph break、recompileを確認する。

graph breakはfusion範囲を切る

Python側のdata-dependent branch、対応していないoperator、動的なside effectなどがあると、compilerが一つのgraphとして扱えずgraph breakが起きる場合がある。

graphが細かく分断されれば、その境界を越えたfusionはできない。さらにcompiled regionとeager regionの間でtensorを受け渡すcostが増える場合もある。

graph breakをゼロにすること自体を目的にしない。可読性や正しさを壊してまで巨大graphへ押し込まず、profileでhot pathを特定して大きな効果がある境界から直す。

dynamic shapeはcompile variantを増やす

LLM servingではbatch size、prompt長、active sequence数が変わる。compilerがshapeごとにspecializeすれば高速なkernelを作りやすいが、新しいshapeのたびにrecompileが発生する。

PyTorchの公式資料では、dynamic shapeを有効にすると再compileを避けるため動的kernelの生成を試みるが、operationや最適化によってspecializationが必要になる場合もある。

shape bucketを使うとvariant数を抑えられる。一方、paddingで余分な計算が増える。compile cache hit率とpadding率を同時に見る。

fusionしすぎると遅くなる理由

大きなfused kernelは、多くの中間値をregisterへ保持する。register使用量が増えると、一つのSMへ同時に置けるwarp数が減り、occupancyが下がる場合がある。registerからlocal memoryへspillすれば、減らしたはずのmemory accessが再び増える。

また、演算ごとに最適なthread blockやtilingが違うこともある。別kernelなら個別に最適化できる処理を無理にまとめると、どちらにも適さない実装になる。

fusion候補は次の条件で判断する。

中間tensorのHBM往復が大きいか
operatorがmemory-boundかcompute-boundか
launch overheadが全体の何割か
fused kernelのregister、shared memory使用量
occupancyとspillが悪化しないか
shapeやdtypeごとの最適実装があるか

PrefillとDecodeで効果が違う

Prefillは大きなmatrix multiplicationやattentionが中心で、一つのkernel実行時間が長い。launch回数削減より、memory trafficとTensor Core利用率の改善が効きやすい。

Decodeは小batchと1-token stepで小さなkernelが増え、launch overheadが目立ちやすい。element-wise fusionやCUDA Graphsの効果が出やすい一方、KV Cache accessがmemory-boundならそのdata layoutも重要になる。

一つの平均latencyで評価せず、PrefillのTTFTとDecodeのITLを分ける。

quantizationではdequantizationの場所が効く

weightをINT8やINT4で保存しても、演算前に大きなtemporary bufferへdequantizeしてHBMへ書けば、memory削減効果を失う。

fused kernelではquantized weightを読み、scaleを適用しながらmatrix multiplicationへ渡せる場合がある。追加のkernel launchとintermediate writeを避けられる。

ただしscale粒度、zero point、group size、activation dtypeがkernelの対応条件と合う必要がある。対応しないshapeで遅いfallbackへ落ちていないか確認する。

数値の正しさを速度と同時に測る

fusionは演算順序やaccumulationの方法を変える場合がある。floating-pointでは加算順序が変わるだけでも最後のbitは変わりうる。低precisionやfast mathを併用すれば差は大きくなる。

比較は一つのlogitだけで行わない。

layerごとの最大絶対誤差と相対誤差
logit、確率分布、top-k tokenの差
greedy decodingの一致率
sampling時の分布と固定seedでの挙動
長文、極端な値、mask境界でのNaN・Inf
downstream taskの品質指標

高速化optionが数学的にexactか、近似や低precisionを含むかを分けて記録する。

compile costとcacheを運用へ含める

JIT compilerは初回にgraph取得、code generation、autotune、compileを行う。起動後の最初のrequestへ載せるとTTFTが跳ねる。

代表shapeをdeployment前にwarmupし、compile cacheを準備する。PyTorchにはcompile artifactをprocess間で再利用する仕組みがあるが、model、code、shape、compiler option、GPU architecture、driver・CUDA versionの互換性を確認する。

cache miss時はserviceを止めずeager pathへfallbackできるようにする。compile時間をrequest deadline内へ隠そうとしない。

custom fused kernelの保守cost

汎用compilerが作れない重要pathでは、TritonやCUDAでcustom kernelを書く選択肢がある。model固有のlayoutやquantizationへ最適化できる一方、保守対象が増える。

GPU architectureごとの性能差
dtype、shape、alignmentの対応範囲
framework APIとcompiler変更
numerical correctness
fallback実装
profilingとdebug symbol

benchmarkで速かった一つのGPUだけを基準にproduction標準へしない。対応外入力を明確にし、安全なreference pathを残す。

監視するmetric

request当たりのkernel launch数
operator別、kernel別の実行時間
HBM read・write量とmemory bandwidth
SM utilization、Tensor Core utilization、occupancy
register、shared memory使用量、spill
fused kernel hit率とfallback率
graph break、recompile、compile cache hit率
compile、autotune、warmup時間
shape、dtype、GPU architecture別の性能
TTFT、ITL、token throughput、p99 latency
reference pathとの数値誤差と品質差

回帰試験

fusionあり・なしで同じ入力のoutputを比較する
batch、sequence length、head dimensionの境界値を試す
causal mask、padding、empty sequenceなどの分岐を通す
FP16、BF16、FP8、INT8・INT4など対応dtypeを分けて試す
対応外shapeが安全にfallbackするか確認する
compile cacheなしのcold startを測る
model・driver・CUDA更新後に再compileと品質確認を行う
concurrent requestでp99 latencyと干渉を見る
NsightなどでHBM traffic、occupancy、spillを確認する
fusionを無効化して即時に切り戻せるか確認する

速さは演算数だけでは決まらない

Kernel Fusionは、連続するGPU演算をまとめ、中間tensorのHBM往復とkernel launch回数を減らす。LLM推論ではnormalization、activation、quantization、attentionなどに使われ、FlashAttentionはIO-aware設計の代表例である。

しかし、大きなkernelはregister pressureやcompile variantを増やし、dynamic shapeではrecompileやfallbackも起きる。fusion数ではなく、HBM traffic、launch数、occupancy、TTFT、ITL、品質差で判断する必要がある。

CUDA Graphsがworkflowのsubmitを効率化するのに対し、Kernel Fusionは一つのkernel内でdata movementを減らす。両方を使う場合も、実trafficのbottleneckをprofileしてから適用することが重要である。