Skip to main content

LLM Inference Optimization

LLM の inference は、parameter サイズと sequence 長の両方で重くなります。Inference 最適化は、production deployment と研究実験の両方で必須です。

何がボトルネックか​

  • Prefill (prompt 処理): GPU compute がボトルネック
  • Decode (token を 1 つずつ生成): GPU メモリ帯域がボトルネック

両者に異なる最適化が必要です。

KV Cache​

Decode 中、過去 token の key / value を再計算しないよう KV cache に貯めます。

step 1: q1 k1 v1
step 2: q2 attends to k1,k2  v1,v2
step 3: q3 attends to k1..k3  v1..v3

KV cache のサイズは

KV cache≈2×L×dhead×nheads×layers×batch\text{KV cache} \approx 2 \times L \times d_{\text{head}} \times n_{\text{heads}} \times \text{layers} \times \text{batch}

で、long context や large batch で支配的なメモリ消費になります。GQA / MQA / MLA はこれを削減する設計です。詳しくは KV Cache を参照してください。

FlashAttention​

FlashAttention は、attention を tile 単位 で SRAM に載せ、softmax を online で計算する高速 attention 実装です。

  • メモリ使用量を O(N)O(N) に削減
  • HBM read/write を削減
  • v2 / v3 でさらに改善

これにより、long context attention が実用速度になります。

PagedAttention と vLLM​

PagedAttention は、KV cache を OS の paging のように固定サイズ block に分割し、論理→物理マッピングで管理する手法です。vLLM の中核技術で、

  • メモリ断片化を解消
  • 高い同時 request 処理 (continuous batching)
  • prefix sharing による cache 共有

を実現します。

Speculative Decoding​

Speculative decoding は、小さい draft model で複数 token を先読みし、大きな target model が一度に検証する手法です。Decode phase の逐次依存を緩和し、target model の forward pass あたりに確定できる token 数を増やします。

Medusa、EAGLE、Lookahead Decoding などの variant を含む詳しい整理は、Speculative Decoding を参照してください。

Quantization​

Weight や activation の精度を下げて memory / compute を削減します。

手法内容
FP16 / BF16標準的な学習・推論
INT8 (LLM.int8, SmoothQuant)Activation outlier を扱う
INT4 (GPTQ, AWQ)4-bit weight quantization
NF4 (QLoRA 推論)学習向け 4-bit format
FP8H100 / Blackwell 世代の native 形式

バッチング戦略​

  • Continuous batching: 異なる長さの request を動的に詰め込む
  • Chunked prefill: 長 prompt を分割し、decode と混在
  • Prefix caching: 同じ system prompt を共有
  • Disaggregated serving: prefill GPU と decode GPU を分離

数式で見る prefill と decode の計算量​

LLM inference は prefill と decode で bottleneck が違います。Prompt 長を TT、hidden dimension を dd、layer 数を LL とすると、prefill の self-attention は概念的に次の計算量を持ちます。

O(LT2d)O(LT^2d)

一方、decode では毎 step の query は 1 token ですが、過去 TT token の KV cache を読む必要があります。

O(LTd)O(LTd)

この式の気持ちは、「prefill は長い prompt をまとめて処理するので compute-bound になりやすく、decode は毎 step 過去の KV を読むため memory bandwidth-bound になりやすい」というものです。

PagedAttention は、KV cache を固定長 block に分け、logical block と physical block の対応を持ちます。Sequence ss の block index を bb とすると、実際の memory address は table lookup で決まります。

addr(s,b)=BlockTable[s,b]\mathrm{addr}(s,b)=\mathrm{BlockTable}[s,b]

この式の気持ちは、「連続した巨大メモリを各 request に確保するのではなく、OS の paging のように小さな block をつなげて KV cache を管理する」というものです。これにより、可変長 request が混在しても memory fragmentation を抑えられます。

関連ページ​

主なソース​