実行モデル
GPUは数千のスレッドを歩調をそろえたグループ（ワープ）で実行し、それらをブロックにまとめてstreaming multiprocessorへ割り当てます。ワープ内のスレッドは1本の命令列を共有するため、分岐が分かれると直列化してレーンを無駄にします。SMあたり十分な数のワープを常駐させること（occupancy）が、メモリレイテンシを隠す手段です。
AI推論ガイド
GPUアーキテクチャとカーネル
サービング基盤のさらに下の層です。モデルのあらゆる演算は最終的にGPUカーネルとして実行され、カーネルの質が、支払ったハードウェアのうち実際に使われる割合を決めます。自分でカーネルを書く機会はまれですが、この層を読むと、その上のシステムがなぜそう振る舞うのかが分かります。
GPUはどう処理を実行するか
メモリ階層
レジスタ、次にオンチップの小さく高速な共有メモリとキャッシュ、そしてオフチップの大きなHBMという順です。オンチップとオフチップの帯域差は桁違いなので、カーネル性能はほぼ、データがその境界をどれだけ渡らずに済むか、そしてアクセスがどれだけ整列するかで決まります。
Tensor Core
行列の積和演算に特化した専用ユニットが、現代のGPUのFLOPSの大半を担います。対応するのは低めの精度（FP16/BF16、FP8、INT8）です。Tensor Coreを使い切れないカーネルや、ハードウェアが持たない精度を使うカーネルは、チップの大部分を遊ばせます。
計算律速かメモリ律速か
どのカーネルも、演算スループットかメモリ帯域のうち先に尽きたほうで頭打ちになります。LLMのprefillは計算が重く、decodeは帯域が重い処理です。ある演算が境界のどちら側にあるかを知れば、FLOPSを増やすべきかデータ移動を減らすべきかが分かります。
カーネル最適化が実際にしていること
- 融合：隣り合う演算（matmul + バイアス + 活性化、attentionの部分ステップ）を1つのカーネルにまとめ、中間結果をHBMへ往復させずオンチップに留めます。カーネルごとの起動オーバーヘッドも消えます。
- タイリング：共有メモリに収まるサイズのブロック単位でデータを処理し、読み込んだタイルを追い出す前に何度も再利用します。
- レイアウトと精度：アクセスが整列するようテンソルを配置し、Tensor Coreがネイティブに実行できる精度を選びます。
- LLMのホットスポット：行列積、attention、正規化、活性化のカーネルが支配的です。サービングエンジンとは、第一近似としては、ひと握りのそうしたカーネルを取り巻くスケジューラです。
カーネルの作業はシステムレベルの取り組みを補完します。continuous batchingやprefix cachingは「どの処理を行うか」を変え、カーネルは「各処理をどれだけ速く実行するか」を変えます。前者についてはサービング最適化をご覧ください。
FlashAttention、その代表例
標準的なattentionはN x Nのスコア行列をHBM上に実体化するため、メモリトラフィックが系列長の二乗で増え、演算は帯域律速になります。FlashAttentionは計算を組み替えます。query、key、valueの各行列をタイル単位でオンチップメモリに流し、タイルが届くたびにsoftmaxの正規化を逐次的に更新し、完全なスコア行列は一度も書き出しません。数学的な結果は同じで、メモリトラフィックはごく一部で済みます。これが長いコンテキストのサービングを実用にしました。
後継バージョンはハードウェア世代ごとに再調整されてきました（作業分割の改善、Tensor Coreの非同期パイプライン、低精度の経路など）。この技術はいまやPyTorchと主要なサービングエンジンすべてに入っています。ここから得られる一般的な教訓は、カーネルの最大の改善は命令の微調整ではなく、メモリ階層を軸に計算を組み替えることから生まれるということです。
ツールの段階
制御の細かさとコストが上がる順に並べています。上から始め、プロファイリングによって上の段では足りないと分かったときだけ下りてください。
- 1. ベンダーライブラリcuBLAS、cuDNNなど、アーキテクチャごとに調整済みの標準演算です。カバー範囲の中では優秀ですが、演算をまたぐ融合はできず、NVIDIA専用です。
- 2. コンパイラtorch.compile、XLA、TVMなど。モデルのグラフを受け取り、融合済みで調整済みのカーネルを自動生成します。標準的なアーキテクチャなら少ない労力で大きな効果が得られますが、目新しい演算は高速経路から外れることがあります。
- 3. カーネルDSLTritonなど。低レベルの詳細はコンパイラに任せつつ、タイル単位のPythonに似た言語で独自カーネルを書けます。フルのCUDA知識なしにattentionの独自派生や融合演算を書ける、ちょうど良い落としどころです。
- 4. 手書きカーネルCUDA（またはROCm）でメモリとワープを明示的に制御します。FlashAttention自身がこの方法で作られています。性能は最大ですが、高い専門性と現実の保守負担、そしてハードウェアへの固定を伴います。
- 5. 移植性を狙うスタック比較的新しいエコシステム（たとえばMojo/MAXやその他のMLIRベースの取り組み）は、ベンダーをまたいで1つのカーネルコードベースを目指します。取り除いてくれる固定と、成熟度とを天秤にかけてください。
判断材料は、自分の演算子がどれだけ標準的か、最大性能にどれだけの価値があるか、どのハードウェアに対応する必要があるか、そしてチームの強みがコンパイラかカーネルか、です。まずプロファイルを取ってください。ボトルネックは直感が指す場所にないことがよくあります。