LLM推論ランタイムをCUDAでゼロから自作して分かったこと — CUDAグラフは「最適化」ではなく「ないと動かない必須部品」だった

7月23日、Anubhab Banerjeeが「How To Build Your Own LLM Runtime From Scratch」と題した記事を公開した。この記事では、NVIDIA H100上でLLM推論ランタイムをゼロからCUDAで自作する実装過程について詳しく紹介されている。

LLMランタイムを自作する意味

llama.cppは速く、安定していて、コントリビュータも多い。本番環境ではそれを使うべきだ、と筆者自身が明言している。ではなぜ自作するのか。

「量子化フォーマットを変えたい、カスタムサンプラーを追加したい、新しいアーキテクチャに対応したい——そのとき、ブラックボックスの中からは何もできない」

デコードスタック全体を自分で握ることで、カーネルを開いて「何が壊れるか」を把握できる状態になる。筆者が公開したのがannotated-llm-runtime（※編集部注：公開時点でのリポジトリURLは記事本文に準拠しているが、大文字・小文字の表記揺れが確認されているため、アクセスできない場合は元記事のリンクを参照されたい）で、約20数ファイルのCUDA/C++ソースで構成され、すべてのホットパスに「なぜそうするのか」のコメントが付いている。

対象モデルはQwen2.5-Coder-7B-Instruct、ターゲットハードウェアはH100（Hopper、sm_90）。

現状のパフォーマンス数値

まず正直な現状値を示す（batch=1、pp512/tg128、warmup 3回）：

メトリクス	本実装	llama.cpp Q4_K_M
TTFT（512トークンプロンプト）	128ms	43ms
デコードITL（定常状態）	16.7 ms/token	4.95 ms/token
デコードスループット	60 tok/s	200 tok/s

llama.cppはコミットb3040、同一GPUクラス・同一モデル・greedy decodeで測定した参照値だ。3倍以上の差がある。（※編集部注：元記事にはH100のSKU——SXM5かPCIeか——の明示がなく、再現時には環境差に注意が必要だ）記事はその差を競うのではなく、「自分のビルドが正常に動いているかの物差し」として使うよう述べている。

最大の収穫：CUDAグラフは「最適化」ではなく「必須」

この記事で最も重要な数字がこれだ。

eager decodeでは1トークンあたり約119ms。CUDAグラフでデコードをキャプチャして再生するだけで約17msに落ちた。7倍の改善、カーネルは同一。

28層×1層あたり約10カーネル、さらにlm_headとargmaxを加えると、1トークン生成に280回以上のカーネルランチが走る。CPUドライバーオーバーヘッドがトークンごとに積み重なり、eager実行では大半の時間がそこに消えていた。

CUDAグラフはカーネルのDAGを一度キャプチャし、以降はその再生のみを行う仕組みで、ドライバーのオーバーヘッドをまとめて除去できる。グラフのキャプチャ自体には制約があり、形状が変わるprefillフェーズには使えない。そのため本実装ではprefillはeager、定常デコードのみグラフ適用という構成を取る。

筆者のコメント：「CUDAグラフはボーナスではない。それがなければランタイムとして機能しない」

Bug #1：`__syncthreads()`がsoftmaxを壊した話

技術的に最も深い内容がこのバグだ。

ページドKVアテンションカーネルはワープ特殊化（warp specialization）を採用している。1つのプロデューサーワープがHopperのTMA（Tensor Memory Accelerator）を使ってHBMからSMEMへKV pageをコピーし、6つのコンシューマーワープがsoftmax＋加重V累算を行う構成だ。

問題は以下のコードにあった：

if (warp_id == 0) {
    // KV pageのTMAコピーを発行
    __syncthreads();  // ← これが致命的なバグ
}

__syncthreads()はブロック全体のバリアである。warp_id == 0の分岐の中にしか置かれていないと、他の6ワープはそのバリアに到達しない。コンシューマーワープはTMAコピーが完了する前にSMEMを読みに行き、古いデータでsoftmaxを計算する。

問題がわかりにくかった理由は、フルページ（16トークン）では「たまたま」コピーが間に合っていたことだ。シーケンス長49のような末尾の端数ページ（最後のブロックに1トークンだけ）になったとき、タイミングがずれてコンシューマーが古いゴミデータを読み、「自信たっぷりに、一貫して、間違った」トークンが出力された。

修正後のコードは以下：

if (warp_id == 0) {
    // ... KV page用TMAを発行 ...
    __syncwarp();   // ← ワープ内32スレッドの同期
    // ... V page用TMAを発行 ...
    __syncwarp();
}
// ブロック全体フェンス：全ワープがここに到達してからSMEMを読む
__syncthreads();

__syncwarp()はワープ内32スレッドの合意、__syncthreads()はブロック全体の合意——この使い分けが肝だ。さらに本実装ではHopperのmbarrier（トランザクション対応バリア）も使い、プロデューサーが転送バイト数を宣言し、コンシューマーがハードウェアレベルで転送完了を待つ仕組みも組み合わせている。

筆者は「通信エンジニアからベアメタルGPUに転向した身として、パケットロストは最終的に再送で解決してきた。CUDAスレッドは再送しない。ただ静かにHopperをヒーターにするだけだ」と述べている。

Hopper特有のGEMV最適化：`prmt.b32`が`__dp4a`を上回る場面

バグ修正と並んで記事が取り上げるのが、GEMV（General Matrix-Vector product）カーネルのHopper固有の最適化だ。

一般にINT4の積和演算には__dp4a（4要素INT8ドット積）が広く使われる。しかしHopperアーキテクチャでbatch=1・小シーケンス長のデコード特有の形状（列方向に極端に細いGEMV）では、prmt.b32命令を使ったビット操作によるアンパック・積算の方が高速になる場合があると筆者は報告している。

prmt.b32はPTXレベルの命令で、32ビットレジスタから任意の4バイトをバイト単位で選択・並び替えてパックできる。INT4量子化ウェイトを効率的にアンパックする際、ワープ内のレジスタ操作だけで完結するため、__dp4aが前提とするバイトアライメント処理のオーバーヘッドを省ける。この最適化の効果はGEMM（行列積）の大きなtileでは薄れ、あくまでデコード時のGEMV形状に限った話だと記事は注記している。

ドキュメントには載っておらず、「実際に書いて計測して初めて分かった」という類の知見だ。

アーキテクチャの概要

本実装の設計方針は一言で言えば「探索空間を意図的に狭くする」ことだ。モデルの定数はすべてconstexprでコンパイル時に固定し、実行時の型分岐を排除している：

static constexpr int kNumHiddenLayers = 28;
static constexpr int kHiddenSize = 3584;
static constexpr int kIntermediateSize = 18944;
static constexpr int kVocabSize = 152064;
// GQA: 28クエリヘッドが4 KVヘッドを共有（7:1比）
static constexpr int kNumAttentionHeads = 28;
static constexpr int kNumKvHeads = 4;
static constexpr int kHeadDim = 128;

重みは対称グループごとINT4（グループサイズ128）でパック。q_projからdown_projまでの7種の線形層に適用し、embed_tokens・RMSNorm・lm_headはFP16のまま残す。「数値バグが出たとき、探索空間を小さく保つための意図的な設計」と説明されている。

重みファイルは.nanoqwen形式で、先頭8バイトがNANOQWENというマジックバイト。これが一致しなければ6GBのデータをVRAMに転送する前に即座に拒否する。ホットパスにYAMLパーサーは存在しない。

まとめ

本記事が示すのは、LLM推論ランタイムを自作することで見えてくる具体的な学びだ。CUDAグラフが「あると速い」ではなく「ないと動かない」という認識の転換、warp特殊化カーネルにおける__syncthreads()の誤用が起こす静かなデータ破壊、そしてprmt.b32が__dp4aを上回るHopper特有のGEMV形状——これらはドキュメントには載っておらず、実際に書いて壊して直した記録だ。

詳細はHow To Build Your Own LLM Runtime From Scratchを参照していただきたい。