필사 모드: カーネルを書く三つの層 — CUDA C++、Triton、CUTLASSを同じ問題で比較する

• はじめに — 同じカーネルを二度書くと見えてくるもの
• 三つの層はそれぞれ何を代わりにやってくれるのか
• CUDA C++ — ハードウェアに直接指示する層
• Triton — タイル単位で書き、スレッドはコンパイラに任せる
  • コンパイルパイプライン — コンパイラが実際にやっていること
  • 何ができたのか見たいとき
  • Tritonが苦手なこと
• 同じカーネルを二度書いて比較する
  • なぜTritonがカスタムアテンションの事実上の標準になったのか
• CUTLASSとCuTe — テンプレートで組み立てるGEMM
  • CuTeとレイアウト代数
• どの層を選ぶか
• 終わりに — 抽象化の層は性能ではなく修正速度で選ぶもの
• 参考資料

はじめに — 同じカーネルを二度書くと見えてくるもの

行単位のsoftmaxカーネルをCUDA C++で書くと、ワープシャッフルによるリダクション、共有メモリでの部分和、2回のブロック同期まで含めて80行ほどになります。同じカーネルをTritonで書くと12行です。それでいて性能は同程度か、Tritonの方がわずかに上回ります。

これを最初に確認すると、「それならCUDA C++を使う理由は何なのか」と思うようになります。そして数か月後には、Tritonではどうしても表現できないカーネルに出会い、またCUDA C++を開くことになります。その二点のどこかに、実務の判断基準があります。

この記事では、カーネルを書く三つの層を同じ物差しで測ります。確認したバージョンはTriton 3.7.1(2026年6月リリース)、CUTLASS 4.6.1(2026年7月)、CUDA Toolkit 13.3 Update 1です。三つの層はいずれも活発に変化し続けているため、APIの詳細は自分が使っているバージョンのドキュメントと照らし合わせてください。

三つの層はそれぞれ何を代わりにやってくれるのか

カーネルを書くという作業は、突き詰めれば六つの決定を下すことです。三つの層の違いは、その決定のうちいくつを自分で行い、いくつを道具に任せるかにあります。

読み方はこうです。CUDA C++は六つすべてを自分で決めます。だから最も自由で、最も時間がかかります。Tritonは最初の二つだけを決め、残りをコンパイラに渡します。だから短く書けますが、その分コンパイラが苦手なパターンでは打てる手が少なくなります。CUTLASSも六つすべてを決めますが、自分の手で書く代わりに、すでに検証済みの部品を選んで組み立てます。だからGEMM系では最高の性能が出ますが、習得には時間がかかります。

CUDA C++ — ハードウェアに直接指示する層

まず基準線です。行単位のsoftmaxを手で書きます。ブロック一つが行一つを担当し、最大値と合計をそれぞれブロック全体でリダクションします。

// softmax_cuda.cu
// ビルド: nvcc -O3 -arch=sm_80 softmax_cuda.cu -o softmax_cuda
#include <cstdio>
#include <cfloat>
#include <cuda_runtime.h>

static const int WARP = 32;
static const int BLOCK = 256;             // ブロックあたり8ワープ

__inline__ __device__ float warpMax(float v) {
  for (int off = WARP / 2; off > 0; off >>= 1)
    v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off));
  return v;
}
__inline__ __device__ float warpSum(float v) {
  for (int off = WARP / 2; off > 0; off >>= 1)
    v += __shfl_xor_sync(0xffffffffu, v, off);
  return v;
}

__global__ void softmaxKernel(float *out, const float *in, int nCols) {
  __shared__ float part[BLOCK / WARP];
  __shared__ float bcast;                 // ブロードキャスト専用スロット

  const long row = blockIdx.x;
  const float *src = in + row * nCols;
  float *dst = out + row * nCols;

  const int tid = threadIdx.x;
  const int lane = tid % WARP, wid = tid / WARP;
  const int nWarps = BLOCK / WARP;

  // パス1: 行の最大値
  float m = -FLT_MAX;
  for (int i = tid; i < nCols; i += BLOCK) m = fmaxf(m, src[i]);
  m = warpMax(m);
  if (lane == 0) part[wid] = m;
  __syncthreads();
  if (tid == 0) {
    float t = part[0];
    for (int w = 1; w < nWarps; w++) t = fmaxf(t, part[w]);
    bcast = t;
  }
  __syncthreads();
  m = bcast;
  __syncthreads();                        // part[]を再利用する前に必須

  // パス2: expの合計
  float s = 0.f;
  for (int i = tid; i < nCols; i += BLOCK) s += __expf(src[i] - m);
  s = warpSum(s);
  if (lane == 0) part[wid] = s;
  __syncthreads();
  if (tid == 0) {
    float t = 0.f;
    for (int w = 0; w < nWarps; w++) t += part[w];
    bcast = t;
  }
  __syncthreads();
  const float inv = 1.f / bcast;

  // パス3: 正規化して書き込み
  for (int i = tid; i < nCols; i += BLOCK) dst[i] = __expf(src[i] - m) * inv;
}

int main() {
  const int R = 8192, C = 4096;
  size_t bytes = (size_t)R * C * sizeof(float);
  float *h = (float *)malloc(bytes);
  for (long i = 0; i < (long)R * C; i++) h[i] = (float)((i * 37) % 100) * 0.01f;

  float *d_in, *d_out;
  cudaMalloc(&d_in, bytes); cudaMalloc(&d_out, bytes);
  cudaMemcpy(d_in, h, bytes, cudaMemcpyHostToDevice);

  for (int i = 0; i < 5; i++) softmaxKernel<<<R, BLOCK>>>(d_out, d_in, C);
  cudaDeviceSynchronize();

  cudaEvent_t a, b; cudaEventCreate(&a); cudaEventCreate(&b);
  cudaEventRecord(a);
  for (int i = 0; i < 50; i++) softmaxKernel<<<R, BLOCK>>>(d_out, d_in, C);
  cudaEventRecord(b); cudaEventSynchronize(b);

  float ms; cudaEventElapsedTime(&ms, a, b);
  double per = ms / 50.0;
  // 最小トラフィック: 1回読んで1回書く。(3パスだがキャッシュが吸収する)
  printf("cuda   %7.3f ms  %7.1f GB/s\n", per,
         2.0 * bytes / (per * 1e-3) / 1e9);
  return 0;
}

80行のうち、実際の数学は3行だけです。残りは全部「ブロックの中で値を集めて、また配り直す」ための配管です。そして、この配管には静かな罠があります。

上のコードにある、m = bcast;の直後の__syncthreads()がそれです。これがないと、速いワープが2パス目のpart[wid]を先に上書きしてしまい、遅いワープがまだ読み取っていない1パス目の結果を消してしまいます。結果はほとんどの実行では正しく出て、たまに間違います。小さい入力ではほとんど再現しません。手書きのカーネルの中で最も高くつく類のバグであり、この配管が長くなるほど発生確率は上がります。

Triton — タイル単位で書き、スレッドはコンパイラに任せる

同じ演算をTritonで書きます。

# softmax_triton.py
# 実行: python softmax_triton.py   (triton 3.7.x、torch 2.13基準)
import torch
import triton
import triton.language as tl


@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_stride, out_stride, n_cols,
                   BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols

    # 行を1本まるごとSRAMに載せる。パディング部分は-infなのでmax/sumに影響しない。
    x = tl.load(in_ptr + row * in_stride + cols, mask=mask, other=-float("inf"))
    x = x - tl.max(x, axis=0)
    num = tl.exp(x)
    y = num / tl.sum(num, axis=0)
    tl.store(out_ptr + row * out_stride + cols, y, mask=mask)


def softmax(x: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = x.shape
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16)
    out = torch.empty_like(x)
    softmax_kernel[(n_rows,)](
        out, x, x.stride(0), out.stride(0), n_cols,
        BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps,
    )
    return out


if __name__ == "__main__":
    torch.manual_seed(0)
    x = torch.randn(8192, 4096, device="cuda", dtype=torch.float32)

    ours, ref = softmax(x), torch.softmax(x, axis=1)
    assert torch.allclose(ours, ref, atol=1e-5), "正確性検証失敗"

    gb = 2 * x.numel() * x.element_size() / 1e9
    for name, fn in [
        ("triton", lambda: softmax(x)),
        ("torch ", lambda: torch.softmax(x, axis=1)),
    ]:
        ms = triton.testing.do_bench(fn, warmup=25, rep=100)
        print(f"{name}  {ms:7.3f} ms  {gb / (ms * 1e-3):7.1f} GB/s")

核心的な違いは、tl.max(x, axis=0)のたった1行です。CUDAでは40行必要だったブロック全体のリダクションが、ここでは配列演算1つで済みます。ワープシャッフルを使うか、共有メモリを何バイト確保するか、同期をどこに入れるか — そのすべてをコンパイラが決めます。そして、先ほど述べた同期漏れのバグは、構造的に発生し得ません。私たちが同期を書くこと自体がないからです。

コンパイルパイプライン — コンパイラが実際にやっていること

Tritonが魔法を使っているわけではありません。段階があります。リポジトリの構造から確認できる経路は次の通りです。

Python関数(@triton.jit)
   │  Python ASTを走査してIRを生成
   ▼
Triton IR(TTIR)          タイル演算。まだハードウェアの概念はない
   │  レイアウト割り当て、コアレッシング、パイプライニングパス
   ▼
TritonGPU IR(TTGIR)      ワープ/スレッド配置と共有メモリが決まる
   │
   ▼
LLVM IR
   │  NVPTXバックエンド        │  AMDGPUバックエンド
   ▼                        ▼
PTX → (ptxas) → SASS       AMDGCN → オブジェクト

ここから、実務上意味のある事実が二つ出てきます。

一つ目は、バックエンドが分岐するのはLLVM IRより下という点です。Tritonリポジトリのthird_partyには、nvidiaとamdのバックエンドが並んで入っています。そのため、同じ@triton.jitカーネルが両方のベンダー向けにコンパイルされます。この性質が、次の二節の話を支えています。

二つ目は、TTIRからTTGIRへ下がる段階が性能の大部分を決定するという点です。レイアウト割り当てを誤ると共有メモリへのアクセスがバンクコンフリクトを起こし、パイプライニングパスが付かないとロード待ちがそのまま露出します。私たちが実際に触れる手がかりはBLOCK_SIZE、num_warps、num_stagesくらいで、そのためTritonのチューニングはたいていこの三つを動かす作業になります。

Tritonは、その調整作業をデコレーターとして提供しています。

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64},  num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64},  num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=4),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 128}, num_warps=4, num_stages=4),
    ],
    key=["M", "N", "K"],          # これらの値が変わるたびに再チューニングする
)
@triton.jit
def matmul_kernel(...):
    ...

keyに書かれた引数が変わるたびに、候補をすべてコンパイルして実測し、最も速いものをキャッシュします。同じことをCUDA C++でやろうとすると、テンプレートのインスタンス化とベンチマークハーネスを自分で書く必要があります。

何ができたのか見たいとき

Tritonをブラックボックスのままにしない方法があります。中間生成物をそのまま取り出せます。

import triton, torch

# カーネルを一度コンパイルしてハンドルを得る
k = softmax_kernel.warmup(
    torch.empty(1, 8, device="cuda"), torch.empty(1, 8, device="cuda"),
    8, 8, 8, BLOCK_SIZE=8, grid=(1,),
)
k._init_handles()

print(k.asm.keys())        # ttir, ttgir, llir, ptx, cubin
print(k.asm["ttgir"][:800])  # レイアウト割り当ての結果を目で確認
print(k.n_regs, k.n_spills)  # レジスタスピルが起きるとここが0でなくなる

n_spillsが0でなければ、レジスタが足りずにローカルメモリに溢れたということで、たいていはBLOCK_SIZEが大きすぎます。Tritonのカーネルが理由もなく遅いときに、まず確認すべき値です。

Tritonが苦手なこと

正直に整理しておくべき部分です。Tritonに頼るべきではない場面は、はっきりと存在します。

• ワープ単位の精密な制御が必要なカーネル。特定のレーン間のシャッフル、ワープ特化(プロデューサーワープとコンシューマーワープを分ける構造)、__ballot系の投票演算に依存するアルゴリズムは、タイル抽象化には収まりません。この隙間を埋めるため、Tritonは最近triton.experimental.gluonというより低い層の言語を実験していますが、3.7.1の時点でもまだ実験段階であり、プロダクションの判断根拠にするにはまだ早いです。
• 不規則なインデックスとデータ依存の制御フロー。グラフ探索、ソート、動的サイズの疎演算のように、アクセスパターンが実行時にしか決まらない場合、タイルモデルは合いません。
• 最高性能を狙ったGEMM。密な行列積はCUTLASSとcuBLASが数年分のチューニングを積み上げてきた領域で、Tritonでその最後の10から20パーセントに追いつくのは困難です。実務上、そうする理由もありません。
• コンパイル時間が応答時間に含まれる場合。JITコンパイルは最初の呼び出しに乗ります。オートチューニングを有効にすると、候補の数だけこれが掛け算になります。サービング経路では、キャッシュのウォームアップが前提条件です。
• 厳密な数値再現性が求められる場合。コンパイラがリダクションの順序を決めるため、バージョンが変わると最後のビットが変わることがあります。

同じカーネルを二度書いて比較する

二つの実装を同じ条件で測る方法と、そのとき出てくる結果の形です。

測定のルールは前の記事と同じです。ウォームアップを捨て、何度も走らせて平均を取り、まず正確性を検証し、時間ではなく実効帯域幅に換算します。Triton側はtriton.testing.do_benchがウォームアップと反復、L2キャッシュのフラッシュまで処理してくれるので、自分で書く必要はありません。CUDA側はcudaEventで同じ条件を揃えます。

softmaxは算術強度の低いメモリバウンドのカーネルなので、どちらの実装も最終的にはHBM帯域幅に突き当たります。そのため、結果の形はおおむね次のようになります。

絶対値はGPUと行の長さによって大きく変わります。上のハーネスを自分のマシンで走らせ、自分の数字を作ってください。この表から読み取るべきは順位ではなく、二つの観察です。

一つ目は、コード量が4倍近く違うのに、性能は同程度だという点です。メモリバウンドのカーネルでは、どちらも同じ壁にぶつかるからです。壁が性能を決めるなら、壁までのコードは短いほど良いということになります。

二つ目は、torch.softmaxがすでにそれだけ速いという点です。これがこの比較の本当の教訓です。標準的な演算子は自分で書かないでください。カーネルを自分で書く価値は、ライブラリにない演算 — つまり融合された組み合わせや、ドメイン特有のマスキングから生まれます。例えば、softmaxの前後にマスキングとスケーリングとドロップアウトが付いていて、それぞれが別々のカーネルとして実行されているなら、一つにまとめた瞬間にHBMの往復が4回から1回に減ります。そのときの利得は3倍以上であり、それこそがTritonを使う理由です。

なぜTritonがカスタムアテンションの事実上の標準になったのか

今、新しいアテンションの変種カーネルが登場すると、その大半はまずTritonで出てきます。理由は性能ではなく構造です。

• アテンションはタイルモデルとよく合います。クエリブロックを一つ取り、キーとバリューをブロック単位で走査しながらオンラインsoftmaxで累積していく構造は、そのままTritonのプログラミングモデルです。
• バリエーションが多く、寿命が短いです。スライディングウィンドウ、ALiBi、ソフトキャッピング、さまざまなスパースパターンが次々に登場し、そのうち一部だけが生き残ります。一つの実験にCUDA C++で2週間もかけるわけにはいきません。
• 一度書けば両方のベンダーで動きます。先ほど見た通り、バックエンドはLLVM IRより下で分岐するため、同じカーネルがNVIDIAとAMDの両方でコンパイルされます。CUDA C++で書くと、AMD側は別作業になります。
• エコシステムがすでにTritonを前提にしています。PyTorchのTorchInductorは、GPU向けコードとしてTritonカーネルを生成します。vLLMはアテンションバックエンドの一覧にTRITON_ATTNを置き、優先順位表にも載せています(vLLMアテンションバックエンドのドキュメント)。すでにインストールされていて、すでに動いているツールだという事実そのものが大きな利点です。

CUTLASSとCuTe — テンプレートで組み立てるGEMM

三つ目の層です。CUTLASSは2017年から続くCUDA C++のテンプレートライブラリで、行列積とその周辺の演算を階層的に分解し、再構成可能な部品にまとめたものです。

核心的な発想は、GEMMを一つの塊として見ないことです。全体の問題をスレッドブロックタイルに分割し、それをさらにワープタイルに、さらにTensor Core命令一つが処理するサイズにまで分割します。各層が独立したテンプレートパラメータであり、私たちはその組み合わせを選びます。

// CUTLASS 3.x/4.xスタイルのGEMM構成の骨格
// 実際にビルド可能な完全な例はリポジトリのexamplesディレクトリを参照してください。
using ElementA = cutlass::half_t;
using ElementB = cutlass::half_t;
using ElementC = float;

// どのサイズのタイルを、何段階のパイプラインで処理するか
using TileShape   = cute::Shape<cute::_128, cute::_128, cute::_64>;
using ClusterShape = cute::Shape<cute::_1, cute::_1, cute::_1>;

// メインループ(データ移動 + Tensor Core累算)と
// エピローグ(結果に付随する後処理)を別々に組み立てる
using CollectiveMainloop = /* CollectiveBuilder<...> */;
using CollectiveEpilogue = /* CollectiveBuilder<...> */;

using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
    cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

ここで実務上最も価値のある概念がエピローグです。GEMMの結果がまだ累算レジスタにある間に、バイアスの加算と活性化関数を適用してしまえば、結果をHBMに書いてからまた読み直す往復がまるごと消えます。cuBLASを呼び出し、別のカーネルで活性化を適用すると、この往復はそのまま残ります。CUTLASSを直接使う理由の大部分はここにあります。

CuTeとレイアウト代数

CUTLASS 3.x以降の基盤がCuTeです。発想は、テンソルのレイアウト — つまり論理座標から物理アドレスへの写像 — を第一級の値にして、それを代数的に組み合わせられるようにすることです。

前の記事では、共有メモリのバンクコンフリクトを配列1マス分のパディングで解決しました。CuTeでは、そうした操作は「スウィズル」というレイアウト変換として表現され、タイルへの分解と合成が可能になり、コンパイル時に検証されます。インデックス演算を手で書いていて間違えるタイプのバグが、型のレベルで捕まります。

学習コストは、正直に言って高いです。テンプレートのエラーメッセージは長く、概念は何層にも重なっています。そのためCUTLASS 4は、PythonインターフェースであるCuTe DSLを追加しました。CUTLASS 4.6.1のREADMEはこれを「C++なしでCuTeの中核概念(レイアウト、テンソル、ハードウェアアトム、スレッドとデータ階層に対する完全な制御)を公開する低レベルのプログラミングモデル」と説明しており、現在は公開ベータであると明記しています。PyPIからnvidia-cutlass-dslとしてインストールでき、4.6.1が最新です。

ベータであるという点は、そのまま素直に受け止めるのがよいでしょう。プロトタイピングには十分使えますが、プロダクションカーネルの唯一の実装にするにはまだ早いです。

どの層を選ぶか

判断を表に圧縮するとこうなります。

実務上の順序として書き直すと、こうなります。まずライブラリでできるかを見る。できなければTritonで書く。Tritonで表現できない、あるいは性能が出ない場合に初めて、プロファイラが指し示すそのカーネルだけをCUDA C++まで下ろす。最初からCUDA C++で始めるのは、そのカーネルがワープ単位の制御を本質的に要求すると確信できているときだけです。

一つ付け加えると、層を混ぜて使うのが正常な状態です。実際の推論スタックを一つ開いてみると、GEMMはcuBLASかCUTLASS、アテンションは手書きのCUDA C++またはTriton、正規化と活性化はTriton、その間をつなぐのはPyTorchというのが普通です。一つを選んで全部を統一しなければならないという圧力に、根拠はありません。

終わりに — 抽象化の層は性能ではなく修正速度で選ぶもの

三つの層を比較しましたが、結論は「Tritonが勝った」ではありません。メモリバウンドのカーネルでCUDA C++とTritonの性能が同程度になったのは、どちらも同じ物理的な壁に突き当たったからであり、GEMMでCUTLASSが前に出るのは、その壁が別の場所にあるからです。層そのものが性能を決めるわけではありません。

層が決めるのは、別のことです。カーネルを一つ直して再測定するまでにかかる時間、同期を漏らして非決定的に間違える確率、新しいGPUが出たときに手を入れ直さなければならないコードの量。性能が同程度であれば、こちらが選択の基準になるべきであり、実務のほとんどでは性能は同程度です。

だから判断基準を一行にまとめると、こうなります。このカーネルを、この先何回くらい修正することになりそうか。一度書いて2年間触らないカーネルなら、CUDA C++まで下りて最後の10パーセントを絞り出すのは合理的です。来月も実験が続くカーネルなら、12行で書いて、その分の時間で実験をもう一回回す方が、ほぼ常に良い選択です。

参考資料

• Tritonリポジトリとチュートリアル: https://github.com/triton-lang/triton
• Triton公式ドキュメント: https://triton-lang.org/main/index.html
• Tillet et al., Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations (MAPL 2019): https://dl.acm.org/doi/10.1145/3315508.3329973
• CUTLASSリポジトリ: https://github.com/NVIDIA/cutlass
• CuTeドキュメント: https://docs.nvidia.com/cutlass/media/docs/cpp/cute/00_quickstart.html
• CUDA C++ Programming Guide: https://docs.nvidia.com/cuda/cuda-c-programming-guide/
• vLLMアテンションバックエンドのドキュメント: https://docs.vllm.ai/en/latest/design/attention_backends.html