C++ と CUDA を用いた高性能 LLM 推論エンジン「tiny-vllm」の構築

このプロジェクトは、軽量で新しい vLLM の妹にあたる「tiny-vllm」です。理論と数学の基礎から、自分で導き出すことを目指しながら学習を進めます。

1. プロジェクトの概要

このリポジトリには以下の 2 つが含まれます：

• 推論サーバーの完全なソースコード
• エンジンの実装プロセスを主導するコース（学習パスツールおよび大学授業資源としても活用可能）

実装対象機能（構成要素）

LLM の本質である行列演算と、GPU 上での最適化を実装します。

• モデル読み込み:

  Safetensors

  形式から LLM モデルの読み込み（例：Llama 3.2 1B Instruct）。
• 完全なフォワードパス: プレフィル（全体入力の処理）とデコード（逐次推論）の実装。
• 計算コア: CUDA カーネルを使用した全計算実装。
• メモリ管理: KV キャッシュ、スタティック・バッチング、クロニウス・バッチング。
• 高度な最適化: オンライン ソフトマックス、FlashAttention に似た実装、PagedAttention。

2. 背景と目的

なぜ C++ と CUDA か？

推論サーバーを構築する際の主要な課題は、ハードウェアの効率化です。

• パフォーマンス最大化: レスポンスタイムの短縮および同時複数プロンプト処理能力の向上。
• GPU の特性活用: LLM 内部の演算（多数の数字の掛け算・足し算）は GPU で高速に実行可能です。
• 数学的本质: LLM の計算は行列乗算であり、線形代数の基礎知識で理解可能です。

AI と知能の視点

• 知能は「大量のパラメータ」と「膨大な計算量」から生まれます。
• モデルの一部を変更し、精度と複雑さのトレードオフを確認できます（例：注意機構の数学的代替案への挑戦）。
• トレーニングフェーズは本コースで扱いません。既訓練モデルの読み込みと高速推論に焦点を当てます。

注記: LLM の設計やゼロからトレーニングする場合は、Karpathy 氏のリポジトリや

tinygrad

（George Hotz）、

micrograd

（Andrei Karpathy）などへの参酌をお勧めします。AI/ML の道で迷った際は、Discord の GPU MODE や Jeremy Howard/Rachel Thomas の書籍も有用です。

3. 環境と事前知識

開発環境

ビルド・実行可能な構成例（NVIDIA GPU 前提）：

• OS: Linux (6.19.8 x64_64)
• CUDA Toolkit: 13.1
• C++ バージョン: C++ 17
• コンパイラ: GCC 15.2.1
• CPU: AMD Ryzen 7 9800X3D
• GPU: NVIDIA RTX 5090

外部依存関係

唯一的外部ライブラリは JSON パーサーです。

#include "json.hpp" // nlohmann/json 3.12.0

構築・実行手順:

1. 依存関係のインストール。

2. ./test.sh

   でプログラムを実行（ビルド直後）。

モデル準備: Hugging Face から

Llama 3.2 1B Instruct

をダウンロードし、リポジトリ内の

model.safetensors

ファイルを使用します。

4. データ型と浮動小数点

BF16 (bfloat16) の採用理由

モデル重量は BF16 で保存されています。推論サーバーの設計ではこれを前提としています。

• Float16 vs BFloat16:
  • Float16: 指数部が狭く、値の範囲（Range）に限界があります。
  • BFloat16: Float32 と同サイズの指数部を持つため、オーバーフロー・アンダーフローを避けやすく、推論タスクに適しています。

浮動小数点の構造 (IEEE 754-2008)

コンピュータ内の実数は、以下の構成ビットで表現されます（16 ビット）。

[ sign | exponent | fraction ]

数値化の仕組み

• Sign (符号): 1 ビット（0:正, 1:負）
• Exponent (指数): 制御範囲を決定。バイアス値が必要。
  • Float16: バイアス 15
  • BFloat16: 8 ビットの指数部
• Fraction (桁): 小数点後の桁数。

数式

$$ (-1)^{\text{sign}} \times 2^{\text{exponent} - \text{bias}} \times (1.\text{fraction}) $$

暗黙の 1: 分数部の前に「1.」が存在しますが、メモリには明示的に保存せず、精度向上のため設計上「暗黙」として扱います。

5. GPU メモリ管理とデータ転送

AI プログラミングでは、Host (CPU) と Device (GPU) の区別が最重要です。

メモリの種類

• Host: PC の DRAM（大容量・低速）。
• Device: GPU の VRAM/HBM（大容量・低速）＋ SRAM（高速・共有メモリ用）。

典型的な処理フロー

1. CPU: 変数宣言・初期化
2. Host メモリサイズ計算 (型サイズ × 要素数)
3. Device へ割り当て (cudaMalloc)
4. データ転送 (cudaMemcpy HostToDevice)
5. カーネル実行
6. 結果転送 (cudaMemcpy DeviceToHost)

コード例：アクティブなトークンの管理

// 1. CPU サイドで入力データを保持
std::vector<int> active_tokens;
active_tokens.push_back(token);

// 2. GPU メモリ確保（最大サイズを想定）
int *gpu_active_tokens = nullptr;
cudaMalloc(&gpu_active_tokens, BATCH_SIZE * sizeof(int));

// 3. データ転送（実際の使用数のみコピー）
int num_active_slots = active_tokens.size();
cudaMemcpy(gpu_active_tokens, 
           active_tokens.data(), 
           num_active_slots * sizeof(int), 
           cudaMemcpyHostToDevice);

// 4. カーネル呼び出し
embeddingGatherDecode(gpu_active_tokens, 
                      num_active_slots, 
                      hidden_state, 
                      weights.embed_tokens);

6. 単一トークン推論の実装

プロジェクト初期化とモデル読み込み

Safetensors ヘッダーを読み取り、演算子の形状・オフセットに基づいて GPU メモリを動的に割り当てます。

CUDA ステータス確認ユーティリティ

int checkGPUStatus() {
    int device_count = 0;
    cudaGetDeviceCount(&device_count);
    if (device_count == 0) return -1; // GPU が見つからないエラー

    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, 0);
    
    std::cout << "Device: " << prop.name << "\n";
    std::cout << "Compute capability: " << prop.major << "." << prop.minor << "\n";
    std::cout << "Global memory: " << prop.totalGlobalMem / (1024 * 1024) << " MB\n";
    
    return 0;
}

モデル重みの読み込み（オフセットマッピング）

ヘッダーから各テンソルのサイズとオフセットを読み、CPU ポインタを GPU メモリ上の位置へマッピングします。

• データ型:

  __nv_bfloat16

  を使用。

// 例：K プロジェクションの重みへのアクセス
// model_weights: GPU ヒープ上の全モデル重みの開始アドレス
// offsets.at("..."): ヘッダーから得たオフセット位置
weights.w_k[layer] = (__nv_bfloat16 *)((char *)model_weights + offsets.at("key_tensor_name"));

トークン化 (Tokenization)

テキスト→整数ベクトル（トークン ID）への変換です。今回は Hugging Face の既存 Tokenizer を使用し、出力を CUDA メモリに転送します。

エンベッディング (Embeddings)

各トークン ID を対応する固定サイズ（例：2048 次元）のベクトルに変換します。

CUDA カーネル #1: Embedding Gather

__global__ void embeddingGatherKernel(...) {
    // グローバルインデックス計算
    int workIndex = threadIdx.x + blockIdx.x * 2048; 
    
    // 入力トークンの重みから埋め込みベクトルを格納
    input_embeddings[workIndex] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x];
}

7. トランスフォーマーブロックの演算実装

RMSNorm (Root Mean Square Layer Normalization)

入力ベクトルを正規化し、安定性を高めます。

$$ \text{normalized}_i = \frac{a_i}{\sqrt{\frac{1}{n}\sum a_i^2} + \epsilon} $$

実装のポイント

• 中間計算の型: BF16 の二乗和などを

  float

  にキャストして精度を維持。
• 除算防衛:

  epsilon

  (例：1e-5) を追加し、ゼロ除算を防ぐ。
• 並列リダクション:

  Tree Reduction

  アルゴリズムを使用して、共有メモリ (

  __shared__

  ) 上でブロック内総和を高速計算。

// RMSNorm カード核イメージ（簡易）
__global__ void rmsNormKernel(...) {
    __shared__ float rms_vector[1024];
    
    // 入力から二乗を計算し、共有メモリに格納
    int workIndex = threadIdx.x + blockIdx.x * 2048;
    rms_vector[threadIdx.x] = (float)input[workIndex] * input[workIndex];
    
    // ... Tree Reduction 処理（総和の累積） ...

    // 正規化計算
    float rms_value = sqrtf(rms_vector[0] / 2048.0 + 1.0e-5);
    
    // 出力へコピー
    output[workIndex] = (__nv_bfloat16)(input[workIndex] * norm_weights[threadIdx.x] / rms_value);
}

RoPE (Rotary Positional Embeddings)

トークンの位置情報を暗黙的に埋め込む技術です。Q と K のベクトルに対して角度回転を行います。

残差接続 (Residual Connections)

入力と出力の要素ごとの加算を行い、多層構造での学習安定性を確保します。

__global__ void residualKernel(...) {
    int workIndex = threadIdx.x + blockIdx.x * 2048;
    input[workIndex] += input_embeds[workIndex];
}

cuBLAS と行列演算の転置罠

LLM の核心（Attention, MLP）は行列乗算です。NVIDIA cuBLAS を使用しますが、データ配列順序に注意が必要です。

• PyTorch/HF: 行優先 (Row-Major)
• cuBLAS: 列優先 (Column-Major) で動作

誤って両者を組み合わせてもダメですので、転置フラグ（

CUBLAS_OP_T

// 例：適切な転置フラグを指定して cuBLAS を呼び出す
cublasGemmEx(handle, 
             CUBLAS_OP_T, // A を転置して扱う
             CUBLAS_OP_N, // B は転置せず
             M, K, N, ... );

8. Attention と推論の仕組み

プレフィル (Prefill) vs デコード (Decode)

• プレフィル: プロンプト全体を入力して最初の出力を生成。入力が大きいため計算コストが高い。
• デコード: 逐次的にトークンを生成。K, V を保存した状態（KV Cache）を利用するため高速化が重要。

KV キャッシュ

以前計算された K (Key) と V (Value) ベクトルを保存するバッファです。これを保持することで、新しいトークン処理時の計算コストを削減します。PagedAttention はこのメモリ管理を最適化する技術です。

Grouped Query Attention (GQA)

複数の Q (Query) ヘッドが 1 つの K/V ヘッドを共有する手法。メモリ効率と計算速度のバランスを取るためのテクニックです。

9. その他の関数実装

• SiLU アクティベーション関数: $$ x / (1 + e^{-x}) $$ ReLU よりも平滑な形状を持ち、負の値も扱う。
• Softmax: 最大値を減算して数値的安定性を確保し、指数関数と除算を実装。オンライン版も存在する。
• Causal Mask: 自己回帰モデルにおいて、トークン $i$ が未来（$j > i$）のデータを見ることを防ぐための三角行列形状のマスク（遮蔽）。
• Argmax: 出力ベクトルから最大値を持つインデックス（次トークン ID）を抽出。通常 CPU/標準ライブラリで行う。

まとめ: この「tiny-vllm」プロジェクトでは、C++ と CUDA を用いて、現代の LLM 推論サーバーがどのように機能するかを第一手から理解・実装します。