AIBit - オープンソースプロジェクトの発見 AIBit - オープンソースプロジェクトの発見

C++とCUDAで自作する高性能LLM推論エンジン

C++とCUDAを使用して、safetensorsからPagedAttentionまでをカバーする、高性能LLM推論エンジンをゼロから構築する方法を学びます。

C++とCUDAで自作する高性能LLM推論エンジン

大規模言語モデルにプロンプトを送信したとき、内部で何が起こっているのか考えたことはありますか？その魔法はモデルの重みだけではなく、計算を orchestrate する推論エンジンにあります。この記事では、vLLMのアーキテクチャに触発されつつも学習ツールとして設計された、C++とCUDAによる独自の高性能LLM推論エンジンの構築方法を解説します。

なぜ独自の推論エンジンを構築するのか？

ほとんどの開発者は、Hugging Face TransformersやOpenAIのような高レベルAPIを通じてLLMと対話します。しかし、レイテンシ、スループット、メモリ使用量といった実際のパフォーマンス特性を理解したり、特定のハードウェアに最適化したりするには、より深く掘り下げる必要があります。推論エンジンをゼロから構築することで、以下のことを学べます：

  • LLMが実際にどのように予測を計算するか
  • 行列演算を高速化するCUDAカーネルの役割
  • GPUワークロードのメモリ管理戦略
  • PagedAttentionや連続バッチ処理のような高度な技術

tiny-vllmのアーキテクチャ

tiny-vllmプロジェクトは、Llama 3.2 1B Instruct向けの完全な推論サーバーを実装しています。以下をカバーしています：

  • Safetensors形式からの実際のLLMの読み込み
  • 完全なフォワードパス（プリフィル＋デコード）
  • カスタムCUDAカーネルによるすべての計算
  • KVキャッシュ管理
  • 静的および連続バッチ処理
  • オンラインソフトマックスとFlashAttention風の技術
  • PagedAttention

はじめに：技術的な前提条件

これを自分のマシンで実行するには、以下が必要です：

  • Linux（6.19.8 x86_64でテスト済み）
  • CUDA Toolkit（13.1）
  • C++17コンパイラ（GCC 15.2.1）
  • NVIDIA GPU（RTX 5090でテスト済み）
  • モデルファイル：Llama 3.2 1B Instructのmodel.safetensors

唯一の外部依存関係はnlohmann/json（単一ヘッダー）です。

Safetensors形式の理解

コードを書く前に、モデルの重みがどのように保存されているかを理解する必要があります。Safetensorsファイルには3つのセクションがあります：

  1. ヘッダーサイズ（8バイト）：ヘッダーの長さを示す64ビット符号なし整数
  2. ヘッダー（JSON）：各テンソルのメタデータ（名前、dtype、形状、オフセット）を含む
  3. テンソルデータ：実際の重みの値

ヘッダーサイズの読み取り方法は次のとおりです：

std::ifstream safetensors_file("model.safetensors", std::ios_base::binary);
uint64_t header_size;
safetensors_file.read(reinterpret_cast<char *>(&header_size), 8);

ヘッダーJSONは、データセクション内の各テンソルの開始位置と終了位置を示します。これにより、不要なデータをコピーせずにテンソルを直接GPUメモリにマッピングできます。

モデルアーキテクチャ：Llama 3.2 1B

Hugging Faceのモデルカードから、アーキテクチャを確認できます：

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 2048)
    (layers): ModuleList(
      (0-15): 16 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
          (k_proj): Linear(in_features=2048, out_features=512, bias=False)
          (v_proj): Linear(in_features=2048, out_features=512, bias=False)
          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((2048,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)
)

フォワードパスは次の順序に従います：

  1. 入力テキストをトークン化
  2. 各トークンの埋め込みを取得
  3. 16のトランスフォーマー層ごとに：
    • RMSNorm
    • グループ化クエリアテンション（RoPE付き）
    • 残差接続
    • RMSNorm
    • フィードフォワードネットワーク（SiLU付き）
    • 残差接続
  4. 最終RMSNorm
  5. 語彙サイズへの線形射影
  6. Argmaxで次のトークンを選択

最初のCUDAカーネル：埋め込み収集

埋め込みルックアップを実装しましょう。各入力トークンについて、埋め込みテーブルから2048要素のベクトルを取得する必要があります。

__global__ void embeddingGatherKernel(int *gpu_input_tokens, __nv_bfloat16 *input_embeddings, __nv_bfloat16 *embed_tokens)
{
    int workIndex = threadIdx.x + blockIdx.x * 2048;
    input_embeddings[workIndex] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x];
    input_embeddings[workIndex + 1024] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x + 1024];
}

ここで注意すべき点：ほとんどのGPUはブロックあたり最大1024スレッドをサポートしますが、各埋め込みには2048要素があるため、スレッドあたり2要素を処理します。これにより、必要以上に多くのブロックを起動するのを防ぎます。

並列リダクションを用いたRMSNorm

RMSNormは各埋め込みベクトルを正規化します。式は次のとおりです：

$$\text{RMS}(a) = \sqrt{\frac{1}{n}\sum_{i=1}^{n}a^2_i}$$

これをGPUで効率的に計算するために、共有メモリ内で並列リダクション（ツリーリダクション）を使用します：

__shared__ float rms_vector[1024];
int workIndex = threadIdx.x + blockIdx.x * 2048;
rms_vector[threadIdx.x] = (float)input[workIndex] * (float)input[workIndex] 
                        + (float)input[workIndex + 1024] * (float)input[workIndex + 1024];
__syncthreads();

// ツリーリダクション
for (int i = 1; i < 1024; i = i * 2)
{
    if (threadIdx.x % (i * 2) == 0)
    {
        rms_vector[threadIdx.x] = rms_vector[threadIdx.x] + rms_vector[threadIdx.x + i];
    }
    __syncthreads();
}

if (threadIdx.x == 0)
{
    rms_vector[0] = sqrt(rms_vector[0] / 2048.0 + 1.0e-5);
}
__syncthreads();

output[workIndex] = (__nv_bfloat16)(((float)input[workIndex] / rms_vector[0]) * (float)norm_weights[threadIdx.x]);
output[workIndex + 1024] = (__nv_bfloat16)(((float)input[workIndex + 1024] / rms_vector[0]) * (float)norm_weights[threadIdx.x + 1024]);

cuBLAS転置のトリック

cuBLASは列優先形式の行列を期待しますが、モデルは通常行優先で保存されています。データを変換する代わりに、行列転置の性質を利用できます：

$$C = A \times B^T \implies C^T = B \times A^T$$

cuBLASは行優先データを列優先として読み取るため（実質的に転置）、それに応じてフラグを設定します：

cublasGemmEx(cublas_handle, 
             CUBLAS_OP_T, CUBLAS_OP_N, 
             KV_DIM, num_active_slots, EMBEDDING_LENGTH, 
             &k_proj_alpha, weights.w_k[layer], CUDA_R_16BF, EMBEDDING_LENGTH, 
             rms_norms, CUDA_R_16BF, EMBEDDING_LENGTH, 
             &k_proj_beta, k_proj_batched_buffer, CUDA_R_16BF, KV_DIM, 
             CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);

KVキャッシュが存在する理由

推論中、各トークンからのKとVの射影のみを保持する必要があります。Q射影、アテンションスコア、フィードフォワード出力は各ステップ後に破棄されます。KとVをキャッシュすることで、新しいトークンごとにそれらを再計算する必要がなくなり、大幅なパフォーマンス向上が得られます。

プリフィルとデコード

  • プリフィル：すべての入力トークンを処理して最初の出力トークンを生成します。これは計算バウンドです。なぜなら、一度に多くのトークンを処理するからです。
  • デコード：後続のトークンを一度に1つずつ生成します。これはメモリバウンドです。なぜなら、キャッシュされたKとVの値を取得するからです。

グループ化クエリアテンション（GQA）

GQAでは、複数のクエリヘッドが同じキーとバリューヘッドを共有します。Llama 3.2 1Bでは、4つのクエリヘッドが1つのキー/バリューヘッドを共有します。これにより、品質を大幅に損なうことなくメモリ使用量が削減され、スループットが向上します。

for (int i = 0; i < NUM_Q_HEADS; ++i)
{
    int k_head_idx = i / GQA_Q_TO_K_RATIO; // 4つのQヘッドが同じ1つのKヘッドを使用
    __nv_bfloat16 *q_head = q_proj + i * HEAD_DIM;
    __nv_bfloat16 *k_head = k_proj[layer] + k_head_idx * HEAD_DIM;
    // ... アテンションスコアを計算
}

連続バッチ処理

静的バッチ処理はN個のリクエストを一緒に処理しますが、最も遅いリクエストを待ちます。連続バッチ処理はスロットを使用してこれを解決します：リクエストが完了すると、そのスロットはキューからの新しいリクエストにすぐに再利用されます。これにより、GPUの使用率が最大化されます。

次は？

リポジトリは活発に開発が進められています。今後のセクションには以下が含まれます：

  • オンラインソフトマックスの導出と実装
  • 効率的なメモリ管理のためのPagedAttention
  • ページングされたKVキャッシュ

結論

LLM推論エンジンをゼロから構築することは、現代のAIシステムがどのように動作するかを深く理解するための最良の方法の一つです。GPUプログラミング、メモリ管理、線形代数、システム設計について学びながら、実際のモデルを実行するものを作成できます。

講師の方であれば、これを教育リソースとして自由に使用してください。学習者の方は、リポジトリをフォークして自分のマシンで動作させ、改善点をコントリビュートしてください。

Happy coding!

ソース

jmaczan/tiny-vllm: C++とCUDAで独自の高性能LLM推論エンジンを構築 - vLLMの小型版