C++とCUDAを使ってゼロからLLM推論エンジンを作る

 (andrewkchan.dev)
15 ポイント 投稿者 GN⁺ 2024-12-16 | 1件のコメント | WhatsAppで共有
  • C++とCUDAを使って、ライブラリなしでLLM推論エンジンを構築する方法
  • これによりLLM推論の全スタックを理解し、さまざまな最適化が推論速度に与える影響を実感できる
  • 目標: 単一のCPU + GPUサーバーで 単一バッチの高速推論 ができるようにモデルを実装し、llama.cpp より速いトークン処理速度を達成する

1. LLMアーキテクチャと推論の概要

  • 主要なLLMの多くは、連続するトランスフォーマーブロックを使う同じアーキテクチャに従っている。
  • モデルのロードは、カスタマイズ可能なトランスフォーマーブロッククラスを定義し、それをシーケンスとして構成して safetensors の重みで初期化することを指す。
  • 推論は主に単一バッチで行われ、「デコード段階」が実行時間の大半を占める。

1.1 推論の概要

  • 推論は、与えられたプロンプトトークンをモデルに渡してKVキャッシュを埋めるプリフィル段階と、モデルを繰り返し実行してトークンを生成するデコード段階に分かれる
    • Prefill段階: プロンプトトークンを処理し、KVキャッシュを初期化
    • Decode段階: 1回に1トークンずつ生成
  • KVキャッシュ: 以前のキー/バリューのペアを保存し、過去コンテキストとのアテンション計算を高速化する
  • モデルのフォワードパスは、埋め込みテーブルを使ってトークンIDを埋め込みベクトルにマッピングし、トランスフォーマーブロックのシーケンスを通じて状態を変換する

1.2 ボトルネックとベンチマーク

  • ボトルネック: 現代のハードウェアではメモリ帯域幅が制約要因になる
    • モデル推論では各トークンを生成するたびにモデル全体を読み込む必要があり、演算性能よりもメモリ帯域幅の制約が大きい
  • モデル量子化は推論速度の改善に効果的
  • 理論上の最大トークン処理量はハードウェアごとに異なり、実際の性能はさまざまな推論エンジンで確認できる
  • 理論上の速度上限:
    • AMD EPYC 7702P: 最大 13.6 tok/s（FP16基準）
    • RTX 4090: 最大 67.1 tok/s（FP16基準）
  • ベンチマーク:
    • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
    • calm: GPU 66 tok/s

2. CPUベースの推論

  • CPUでの初期実装はシングルスレッドで行われ、FP32重みのみをサポートする
  • マルチスレッディングでコードの並列化を始め、SIMDを使って性能を向上できる

2.1 マルチスレッディング

  • OpenMPを使って行列-ベクトル積（matmul）とマルチヘッドアテンションを並列化し、性能を改善
  • 最適化結果: 速度が 0.6 tok/s → 4.4 tok/s に改善

2.2 重みの量子化とSIMD最適化

  • 量子化: FP32重みをFP16に量子化し、メモリ使用量を半減して性能を改善
  • SIMD: AVX2を使って8個のFP32値を同時に処理するよう最適化
  • 結果: 8.4 tok/s を達成

3. GPUベースの推論

  • モデルをFP16に量子化してRTX 4090にロードし、GPU推論の実装を始められる
  • CUDAを使うことで、C++関数（カーネル）をGPU上で並列実行できる

3.1 CUDAへの単純移植

  • CPU演算を1対1でCUDAカーネルに変換することでGPUバックエンドを実装できる
  • CUDAカーネルは非同期に実行されるが、同一ストリーム内では逐次実行される
  • 問題点: スレッド効率が悪く、GPUリソースを十分に活用できない → 2.9 tok/s と低速

3.2 より良い行列積（matmul）

  • 行列積はCPUで大きなランタイムを占め、OpenMPで最適化できる
  • GPUではブロックごとに1行を処理するようにして、スレッド活用率を高められる
  • 最適化方法:
    1. 1ブロックが1行を処理し、ブロック内スレッドが協調して計算
    2. warp単位の集約（warp reduction） を適用
  • 結果: 51.7 tok/s に高速化

3.3 カーネル融合と追加最適化

  • カーネルを融合して性能を向上できる
    • カーネル融合: 連続する演算を1つのカーネルにまとめ、メモリアクセスと計算時間を最小化
  • メモリアクセスパターンの最適化と 空間再利用 により 56.1 tok/s を達成

3.4 Attention最適化と長いコンテキスト処理

  • 問題点: 長いコンテキストではアテンションカーネルが性能ボトルネックになる
  • 解決策:
    1. メモリアクセス最適化: 連続したメモリブロックを読み込むように再設計
    2. atomicAdd の代わりに共有メモリを使い、欠落する小数値の問題を解決
  • 最適化結果:
    • 短いコンテキスト: 63.8 tok/s（llama.cppの61.0 tok/sより高速）
    • 長いコンテキスト: 58.8 tok/s を達成

3.5 KVキャッシュ量子化とコンパイラ最適化の問題

  • KVキャッシュをFP16に量子化すると性能低下が発生（コンパイラ最適化不足）
  • 解決策: ループを手動でアンローリングし、メモリプリフェッチを適用
  • 結果: FP32比で約2倍の速度向上 を実現し、長いコンテキスト性能 58.8 tok/s を維持

4. 今後の改善方向

  • プロンプトprefill最適化: 複数トークンを同時に処理して最初のトークン生成時間を短縮
  • Attentionカーネル融合: FlashAttentionのような最適化手法を適用
  • より高い量子化: FP8、INT8、INT4の適用および活性化/キャッシュ量子化
  • カーネル最適化: メモリ帯域幅と演算効率を最大化する高度な手法を導入
  • ライブラリ利用: cuDNN、cuBLASのようなライブラリを活用して最適化時間を短縮

結果の要約:

  • CPUとGPUでのさまざまな最適化により 63.8 tok/s を達成
  • llama.cppやcalmに近い、あるいはそれを上回る性能を記録
  • ライブラリなしで C++とCUDAのみ により高性能なLLM推論エンジンを実装

関連記事

1件のコメント

 
GN⁺ 2024-12-16
Hacker Newsの意見
  • 投稿者は自分のブログ記事が注目を集めてうれしく、フィードバックを聞きたがっている
  • ある読者は記事が素晴らしいと称賛し、執筆にどれくらい時間がかかったのか気にしている
    • GPGPU 分野で働く人として似たような記事を書きたいが、所要時間が読めないためためらっている
  • 別の読者は、コードが tensor cores や wgmma 命令を活用していないと考えている
    • この種のプログラミングは複数の作業を同時に処理しなければならないため難しいと説明している
    • 帯域幅の制約により、追加の演算は必要ないかもしれないと言及している
    • ブログのコードは他のアクセラレータへ移植する際にもうまく機能する可能性が高いと評価している
    • wgmma を使うと Nvidia の世代間移植性が低下するのではないかと懸念している
  • また別の読者は、これに似た Python の資料を探しており、チームと共有したいと考えている
    • 性能よりも、概念的に完全でチュートリアル風に簡潔な資料を求めている
  • あるユーザーは、自分の Mistral バージョンと tokens/second の性能を比較したがっている
    • README の量子化セクションを参照するよう勧められている
  • __shfl_down は最近では warp 同期の問題があるため推奨されないという意見もある