Tiny-vLLMでゼロから構築する高性能LLM推論エンジン

現代の大規模言語モデル（LLM）の複雑さは、それらを実際に動かすために必要な根本的なエンジニアリングをしばしば隠してしまいます。PyTorchのようなハイレベルなライブラリはモデル設計を容易にしますが、重みファイルから応答性の高い推論サーバーへの移行には、メモリ管理、GPUカーネルエンジニアリング、および線形代数への深い理解が必要です。

Tiny-vLLMは、このプロセスを解明するために設計された教育的なプロジェクトです。ブラックボックスを提供するのではなく、Llama 3.2 1B Instructモデルをリファレンスとして使用し、C++とCUDAを用いて高性能な推論エンジンを構築するためのガイド付きコースとして機能します。それは、理論的なアーキテクチャとハードウェア加速の泥臭い現実との間のギャップを埋めるものです。

推論サーバーの解剖学

推論エンジンを構築するには、まずLLMが物理的に何であるかを理解しなければなりません。それは、数百万の浮動小数点数（重み）を含むファイルと、それらの数値がどのように一連の操作で使用されるかを定義する設計図（アーキテクチャ）です。

推論サーバーの主な役割は、その設計図を実行可能なコードに変換することです。高いパフォーマンスを得るために、これは通常、ハードウェアの利用率を最大化するためにC++とCUDAで行われます。主な目的は、レイテンシを最小限に抑えつつ、複数のプロンプトを同時に処理することであり、そのためには計算をCPUからGPUへ移動させる必要があります。GPUでは、LLMの要である行列乗算を数千のコアで並列化できます。

モデルのロード：SafetensorsとBF16

ほとんどの現代的なモデルはSafetensors形式で配布されています。Safetensorsファイルは、ヘッダーサイズ、テンソルメタデータ（dtype、shape、およびオフセット）を含むJSONヘッダー、および生のテンソルデータで構成されています。

Bfloat16の選択

LLM推論における重要な技術的決定は、数値形式です。標準的な16ビット浮動小数点数（FP16）は一般的ですが、多くのモデルは**Bfloat16 (BF16)**を使用します。

BF16は精度と範囲のトレードオフを行います。指数部には8ビット（32ビット浮動小数点数と同じ）を使用し、仮数部にはわずか7ビットを使用します。これにより、トランスフォーマー層で必要とされる大規模な加算中に、オーバーフローやアンダーフローの問題が発生しにくくなります。経験的な証拠によれば、精度の低下は、それが提供する安定性のための許容可能なトレードオフであることが示唆されています。

GPUメモリ管理

効率的な推論には、**Host (CPU/DRAM)とDevice (GPU/VRAM)**の間のデータ移動に関する厳格な戦略が必要です。GPUはシステムのDRAMに直接アクセスできないため、データはcudaMallocを介して明示的に割り当てられ、cudaMemcpyを介して転送されなければなりません。

パフォーマンスを最適化するために、エンジンはいくつかの黄金律に従います：

割り当てを最小限にする: 大きなバッファを事前に割り当てます。
メモリを再利用する: データの使用が終わった後に再利用できるバッファを特定するために、生存期間分析（lifetime analysis）を使用します。
転送を減らす: GPUへのデータコピーを可能な限り少なくします。

CUDAカーネルのエンジニアリング

トランスフォーマーのフォワードパスを実装するには、カスタムCUDAカーネルを記述する必要があります。これには、**SIMT (Single Instruction, Multiple Threads)**の観点から考えることが含まれます。そこでは、同じ関数がブロックとワープにグループ化された数千のスレッドで実行されます。

Embedding Gather

最初のステップは、入力トークン（整数）を対応する埋め込みベクトルにマッピングすることです。ほとんどのGPUは1ブロックあたり1024スレッドの制限がありますが、Llama 3.2の埋め込みは2048要素の長さであるため、一般的な最適化として、各スレッドが埋め込みの2つの要素を処理するようにし、ブロックあたりのスループットを実質的に2倍にします。

RMSNormと並列リダクション

Root Mean Square Layer Normalization (RMSNorm) は、埋め込み全体にわたる二乗値の平均の平方根を計算する必要があります。これは同期の課題を生み出します。つまり、異なるスレッド間でどのように値を合計するか、という課題です。

Tiny-vLLMは、これを**Parallel Reduction (Tree Reduction)**を用いて実装しています。単一の共有変数（これは競合状態を引き起こします）を使用する代わりに、スレッドはローカルな合計を__shared__メモリベクトルに書き込みます。次に、それらをツリー構造のように反復的に合計し、インデックス0に単一の合計値が残るまで、各ステップでアクティブなスレッド数を半分にしていきます。

数値的な安定性を確保し、ゼロ除算（これはNaNを発生させ、推論をクラッシュさせます）を防ぐために、RMSの計算にイプシロン値（例：1e-05）が加えられます。

行列乗算の課題

行列乗算は、エンジンの中で最も計算負荷の高い部分です。カスタムカーネルを記述することもできますが、業界標準はNVIDIAの高度に最適化された線形代数ライブラリであるcuBLASです。

しかし、大きな摩擦点が存在します：cuBLASは行列をcolumn-major format（列優先形式）として期待しますが、LLMの重みは通常row-major format（行優先形式）として保存されています。Tiny-vLLMは、メモリ内のデータを物理的に並べ替えることを避けるために、数学的な転置トリックを使用します。転置フラグ（CUBLAS_OP_TとCUBLAS_OP_N）とGEMM (General Matrix Multiply) コールの順序を操作することで、エンジンは行優先データを列優先として扱うことができ、オーバーヘッドを大幅に削減できます。

スループットの最適化：バッチングとキャッシュ

Prefill vs. Decode

推論は2つの異なるフェーズに分かれます：

Prefill: エンジンは最初のトークンを生成するために、入力プロンプト全体を処理します。これは計算負荷が高いですが、非常に並列化可能です。
Decode: エンジンは後続のトークンを一つずつ生成します。これは、最後に生成されたトークンのみを処理するため、メモリ帯域幅に制限される（memory-bandwidth bound）プロセスです。

KV Cache

デコードフェーズにおいて、以前のすべてのトークンに対するKey (K) と Value (V) の射影を再計算することを避けるために、エンジンはKV Cacheを実装しています。このキャッシュは、処理されたすべてのトークンのKとVベクトルを保存し、、モデルが新しいトークンの射影を単に加算し、既存の履歴に対してアリーション（attention）的操作を行えるようにします。

Advanced Batching

サーバーをスケールさせるために、Tiny-vLLMは2種類のバッチングを探索します：

Static Batching: 複数のリクエストを一度に処理します。これはスループットを向上させますが、レイテンシを導入します。なぜなら、バッチ全体が最も長いシーケンスの終了を待たなければならないからです。
Continuous Batching: より洗練練されたアプローチであり、新しいリクエストがスロットが空いた瞬間にバッチに入ることができ、静的なパディングによる無駄を排除します。

結論

推論エンジンを構築することから、LLMの「知能」は単一のコード行に見出出されるのではなく、パラメータの膨大なスケールと計算の効率性から立ち現上がるものであることがわかります。RMSNormにおけるツリーリダクションや、cuBLASの転置トリックのような低レベルのプリミティブを実装することで、開発者はLLMをサービスとして利用する段階を超え、現代のAIを可能にするハードウェアとソフトウェアの協調設計（hardware-software co-design）を理解し始めることができます。