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を参照していただきたい。