TileLang で GPU カーネル設計を解説

本チュートリアルでは、TVM を介してパフォーマンス指向の GPU カーネルを設計・コンパイルするための高レベル Python ドメイン固有言語である TileLang について探ります。まずは CUDA 環境の検証を行い、再利用可能なベンチマークおよび数値検証ユーティリティを整備します。その後、ベクトル加算からタイル化されたテンソルコア行列乗算、スケジュール探索、融合 GEMM エピローグ、行ごとの Softmax、そして FlashAttention へと段階的に実装を進めていきます。

本チュートリアルを通じて、TileLang の共有メモリタイルやレジスタ断片、パイプライン化されたループ、並列反復プリミティブ、還元演算、テンソルコア GEMM 演算子を直接取り扱います。その一方で、スレッドマッピング、メモリアウトライン、同期処理、ベクトル化、および低レベル CUDA 命令の生成はコンパイラーに任せる形となります。

また、作成したカーネルを PyTorch や cuBLAS のベースラインと比較し、生成された CUDA ソースコードを検証。メモリと計算のスループットを評価するとともに、オートチューニングを用いてアーキテクチャ依存のカーネル設定を特定します。

import os

import sys

import math

import time

import subprocess

import traceback

def _sh(cmd: str) -> int:

print(f"$ {cmd}", flush=True)

return subprocess.run(cmd, shell=True).returncode

def _bootstrap():

"""tilelang がインストールされていない場合は自動でインストールします。まずは安定版の wheel を試し、失敗したら夜間ビルド版を試みます。"""

try:

import tilelang

return

except ImportError:

pass

print(">> installing tilelang (bundled TVM を読み込みます。約 1〜3 分)")

_sh(f"{sys.executable} -m pip install -q tilelang")

try:

import tilelang

return

except ImportError:

print(">> 安定版 wheel が利用できないため、夜間ビルドチャンネルを試します")

_sh(f"{sys.executable} -m pip install -q tilelang "

f"-f https://tile-ai.github.io/whl/nightly")

import tilelang

if os.path.isdir("/usr/local/cuda"):

os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")

os.environ["PATH"] = os.environ.get("PATH", "") + ":/usr/local/cuda/bin"

_bootstrap()

import torch

import torch.nn.functional as F

import tilelang

import tilelang.language as T

def banner(title: str):

print("\n" + "=" * 78)

print(f" {title}")

print("=" * 78, flush=True)

def bench(fn, warmup: int = 10, rep: int = 50) -> float:

"""CUDA イベントを用いて測定した、ミリ秒単位の中央値に近いレイテンシを返します。"""

for _ in range(warmup):

fn()

torch.cuda.synchronize()

start, end = torch.cuda.Event(True), torch.cuda.Event(True)

start.record()

for _ in range(rep):

fn()

end.record()

torch.cuda.synchronize()

return start.elapsed_time(end) / rep

def check(actual: torch.Tensor, ref: torch.Tensor, name: str, tol: float = 2e-2) -> bool:

"""相対フロベニウスノルムによる検証を行います。fp16 の場合、絶対誤差 (atol) よりも意味があります。"""

a, r = actual.float(), ref.float()

rel = (a - r).norm() / r.norm().clamp_min(1e-12)

amax = (a - r).abs().max().item()

ok = bool(rel < tol) and math.isfinite(amax)

print(f" [{'PASS' if ok else 'FAIL'}] {name}: rel_err={rel:.3e} max_abs={amax:.3e}")

return ok

banner("0. ENVIRONMENT")

assert torch.cuda.is_available(), "GPU が利用できません。ランタイムを変更して GPU を選択してください。"

DEV = torch.device("cuda")

PROPS = torch.cuda.get_device_properties(0)

CC = torch.cuda.get_device_capability(0)

SM = CC[0] * 10 + CC[1]

print(f" tilelang : {getattr(tilelang, '__version__', 'unknown')}")

print(f" torch : {torch.__version__}")

print(f" GPU : {PROPS.name} (sm_{SM}, {PROPS.multi_processor_count} 個の SM、{PROPS.total_memory/2**30:.1f} GiB)")

SMEM_CAP = 48 * 1024 if SM < 80 else 96 * 1024

DEFAULT_STAGES = 2 if SM < 80 else 3

print(f" smem budget: {SMEM_CAP//1024} KB/block, default num_stages={DEFAULT_STAGES}")

print(" note: tilelang はコンパイル済みカーネルを ~/.tilelang/cache にキャッシュします。そのため、")

print(" 同じセルを二度実行すると、大幅に高速化されます。")

@tilelang.jit(out_idx=[-1])

def make_vector_add(N: int, block_N: int = 256, dtype: str = "float32"):

@T.prim_func

def main(A: T.Tensor((N,), dtype),

B: T.Tensor((N,), dtype),

C: T.Tensor((N,), dtype)):

with T.Kernel(T.ceildiv(N, block_N), threads=256) as bx:

for i in T.Parallel(block_N):

C[bx * block_N + i] = A[bx * block_N + i] + B[bx * block_N + i]

return main

def section_1():

banner("1. HELLO, TILE — ベクトル加算と生成された CUDA コード")

N = 1 << 22

add = make_vector_add(N, block_N=256)

a = torch.randn(N, device=DEV)

b = torch.randn(N, device=DEV)

c = add(a, b)

check(c, a + b, "vector_add")

ms = bench(lambda: add(a, b))

gbs = 3 * N * 4 / (ms * 1e-3) / 1e9

ref_ms = bench(lambda: a + b)

print(f" tilelang: {ms*1e3:8.1f} us ({gbs:6.1f} GB/s)")

print(f" torch : {ref_ms*1e3:8.1f} us <- どちらも純粋な帯域幅依存なので、引き分けが正解です")

src = add.get_kernel_source()

print("\n --- 生成されたデバイスコード (先頭 32 行) ---")

for line in src.splitlines()[:32]:

print(" " + line)

print(" ---------------------------------------------")

Google Colab の CUDA 環境を設定し、TileLang を nightly バージョンをフォールバック先としてインストールします。その後、必要な PyTorch と TileLang のモジュールを読み込みます。

カーネルのレイテンシを測定し、相対誤差を用いて数値出力を比較する再利用可能なベンチマーク、検証、レポート作成ユーティリティを定義しました。

次に、TileLang によるベクトル加算カーネルを実装して GPU で実行し、その帯域幅を PyTorch と比較します。さらに、コンパイラが生成した CUDA ソースコードも確認します。

@tilelang.jit(out_idx=[-1])

def make_matmul(M: int, N: int, K: int,

block_M: int = 128, block_N: int = 128, block_K: int = 32,

num_stages: int = 3, threads: int = 128,

use_swizzle: bool = False,

dtype: str = "float16", accum_dtype: str = "float"):

@T.prim_func

def main(A: T.Tensor((M, K), dtype),

B: T.Tensor((K, N), dtype),

C: T.Tensor((M, N), dtype)):

with T.Kernel(T.ceildiv(N, block_N),

T.ceildiv(M, block_M),

threads=threads) as (bx, by):

A_shared = T.alloc_shared((block_M, block_K), dtype)

B_shared = T.alloc_shared((block_K, block_N), dtype)

C_local = T.alloc_fragment((block_M, block_N), accum_dtype)

if use_swizzle:

T.use_swizzle(panel_size=10, enable=True)

T.clear(C_local)

for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):

T.copy(A[by * block_M, ko * block_K], A_shared)

T.copy(B[ko * block_K, bx * block_N], B_shared)

T.gemm(A_shared, B_shared, C_local)

T.copy(C_local, C[by * block_M, bx * block_N])

return main

def smem_bytes(block_M, block_N, block_K, stages, itemsize=2):

return (block_M * block_K + block_K * block_N) * itemsize * stages

def section_2():

banner("2. メモリ階層 — タイル化された Tensor Core GEMM")

M = N = K = 2048

bm, bn, bk, st = 128, 128, 32, DEFAULT_STAGES

while smem_bytes(bm, bn, bk, st) > SMEM_CAP and st > 1:

st -= 1

print(f" config: {bm}x{bn}x{bk}, num_stages={st}, "

f"smem={smem_bytes(bm,bn,bk,st)/1024:.0f} KB")

kernel = make_matmul(M, N, K, bm, bn, bk, num_stages=st, threads=128)

a = torch.randn(M, K, device=DEV, dtype=torch.float16)

b = torch.randn(K, N, device=DEV, dtype=torch.float16)

c = kernel(a, b)

check(c, a @ b, "matmul 2048^3")

flops = 2 * M * N * K

ms = bench(lambda: kernel(a, b))

ms_ref = bench(lambda: a @ b)

print(f" tilelang : {ms:7.3f} ms -> {flops/(ms*1e-3)/1e12:6.2f} TFLOP/s")

print(f" cuBLAS : {ms_ref:7.3f} ms -> {flops/(ms_ref*1e-3)/1e12:6.2f} TFLOP/s")

print(f" ratio : {ms_ref/ms*100:5.1f}% of cuBLAS from ~20 lines of Python")

src = kernel.get_kernel_source()

for needle in ("mma.sync", "wgmma", "ldmatrix", "cp.async", "tl::gemm"):

if needle in src:

print(f" emitted: {needle}")

return kernel

def section_3():

banner("3. パラメータ調整 — スケジュールを手動で探索する")

M = N = K = 2048

a = torch.randn(M, K, device=DEV, dtype=torch.float16)

b = torch.randn(K, N, device=DEV, dtype=torch.float16)

flops = 2 * M * N * K

candidates = [

(64, 64, 32, 2, 128, False),

(128, 128, 32, 2, 128, False),

(128, 128, 32, 2, 128, True),

(128, 128, 32, 3, 128, False),

(128, 128, 64, 2, 256, False),

(128, 256, 32, 2, 256, False),

]

print(f" {'tile':>16} {'stg':>4} {'thr':>4} {'swz':>4} {'smem':>7} "

f"{'ms':>8} {'TFLOP/s':>9}")

results = []

for bm, bn, bk, st, thr, swz in candidates:

need = smem_bytes(bm, bn, bk, st)

tag = f"{bm}x{bn}x{bk}"

if need > SMEM_CAP:

print(f" {tag:>16} {st:>4} {thr:>4} {str(swz):>4} {need//1024:>5}KB "

f" SKIPPED (over smem budget)")

continue

try:

k = make_matmul(M, N, K, bm, bn, bk, num_stages=st,

threads=thr, use_swizzle=swz)

c = k(a, b)

ok = (c - (a @ b)).float().norm() / (a @ b).float().norm() < 2e-2

ms = bench(lambda: k(a, b), warmup=5, rep=20)

results.append((ms, tag, st, thr, swz))

print(f" {tag:>16} {st:>4} {thr:>4} {str(swz):>4} {need//1024:>5}KB "

f"{ms:>8.3f} {flops/(ms*1e-3)/1e12:>9.2f}"

f"{'' if ok else ' <- NUMERICALLY WRONG'}")

except Exception as e:

print(f" {tag:>16} {st:>4} {thr:>4} {str(swz):>4} - "

f"failed: {type(e).__name__}")

if results:

best = min(results)

print(f"\n winner: {best[1]}, stages={best[2]}, threads={best[3]}, "

f"swizzle={best[4]} ({best[0]:.3f} ms)")

print(" Takeaway: the best schedule is arch- and shape-dependent, which is")

print(" exactly why section 7 exists.")

入力タイルをグローバルメモリ、共有メモリ、レジスタフラグメント間を移動させるタイル化されたテンソルコア行列乗算カーネルを実装しました。タイルサイズやパイプライン段数、スレッド数、L2 スイッシング制御などを細かく設定しつつ、TileLang にテンソルコア命令の生成、同期処理、メモリ転送ロジックの作成を任せています。

その後、複数のスケジューリング構成でベンチマークを実施し、数値精度を検証。その上で、ハードウェアに依存する最高性能を発揮するカーネル構成を特定しました。

@tilelang.jit(out_idx=[-1])

def make_matmul_bias_gelu(M: int, N: int, K: int,

block_M: int = 128, block_N: int = 128, block_K: int = 32,

num_stages: int = 2, threads: int = 128,

dtype: str = "float16", accum_dtype: str = "float"):

@T.prim_func

def main(A: T.Tensor((M, K), dtype),

B: T.Tensor((K, N), dtype),

Bias: T.Tensor((N,), dtype),

C: T.Tensor((M, N), dtype)):

with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M),

threads=threads) as (bx, by):

A_shared = T.alloc_shared((block_M, block_K), dtype)

B_shared = T.alloc_shared((block_K, block_N), dtype)

Bias_shared = T.alloc_shared((block_N,), dtype)

C_local = T.alloc_fragment((block_M, block_N), accum_dtype)

T.clear(C_local)

T.copy(Bias[bx * block_N], Bias_shared)

for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):

T.copy(A[by * block_M, ko * block_K], A_shared)

T.copy(B[ko * block_K, bx * block_N], B_shared)

T.gemm(A_shared, B_shared, C_local)

for i, j in T.Parallel(block_M, block_N):

C_local[i, j] = C_local[i, j] + Bias_shared[j]

for i, j in T.Parallel(block_M, block_N):

C_local[i, j] = C_local[i, j] / (

1.0 + T.exp(-1.5957691216 * (

C_local[i, j] + 0.044715 * C_local[i, j]

• C_local[i, j] * C_local[i, j])))

T.copy(C_local, C[by * block_M, bx * block_N])

return main

def section_4():

banner("4. EPILOGUE FUSION — GEMM + bias + GELU in one kernel")

M, N, K = 4096, 4096, 1024

st = DEFAULT_STAGES

while smem_bytes(128, 128, 32, st) > SMEM_CAP and st > 1:

st -= 1

fused = make_matmul_bias_gelu(M, N, K, 128, 128, 32, num_stages=st, threads=128)

a = (torch.randn(M, K, device=DEV, dtype=torch.float16) / K ** 0.25)

b = (torch.randn(K, N, device=DEV, dtype=torch.float16) / K ** 0.25)

bias = torch.randn(N, device=DEV, dtype=torch.float16)

out = fused(a, b, bias)

ref = F.gelu(((a @ b).float() + bias.float()), approximate="tanh")

check(out, ref, "matmul+bias+gelu", tol=3e-2)

ms_f = bench(lambda: fused(a, b, bias))

ms_e = bench(lambda: F.gelu(a @ b + bias, approximate="tanh"))

print(f" fused (1 kernel) : {ms_f:7.3f} ms")

print(f" torch (3 kernels) : {ms_e:7.3f} ms")

print(f" speedup : {ms_e/ms_f:5.2f}x")

print(f" HBM traffic saved : ~{2*M*N*2/2**20:.0f} MiB of intermediate reads/writes")

@tilelang.jit(out_idx=[-1])

def make_softmax(M: int, N: int, block_M: int = 4, threads: int = 128,

dtype: str = "float16", accum_dtype: str = "float"):

@T.prim_func

def main(X: T.Tensor((M, N), dtype),

Y: T.Tensor((M, N), dtype)):

with T.Kernel(T.ceildiv(M, block_M), threads=threads) as bx:

X_shared = T.alloc_shared((block_M, N), dtype)

X_local = T.alloc_fragment((block_M, N), accum_dtype)

row_max = T.alloc_fragment((block_M,), accum_dtype)

row_sum = T.alloc_fragment((block_M,), accum_dtype)

T.copy(X[bx * block_M, 0], X_shared)

T.copy(X_shared, X_local)

T.reduce_max(X_local, row_max, dim=1, clear=True)

for i, j in T.Parallel(block_M, N):

X_local[i, j] = T.exp(X_local[i, j] - row_max[i])

T.reduce_sum(X_local, row_sum, dim=1)

for i, j in T.Parallel(block_M, N):

X_local[i, j] = X_local[i, j] / row_sum[i]

T.copy(X_local, Y[bx * block_M, 0])

return main

def section_5():

banner("5. REDUCTIONS — row-wise softmax")

M, N = 8192, 1024

sm = make_softmax(M, N, block_M=4, threads=128)

x = torch.randn(M, N, device=DEV, dtype=torch.float16) * 3.0

y = sm(x)

check(y, torch.softmax(x.float(), dim=-1), "softmax")

ms = bench(lambda: sm(x))

ms_ref = bench(lambda: torch.softmax(x, dim=-1))

gbs = 2 * M * N * 2 / (ms * 1e-3) / 1e9

print(f" tilelang: {ms*1e3:7.1f} us ({gbs:6.1f} GB/s)")

print(f" torch : {ms_ref*1e3:7.1f} us")

print(" Both are memory bound; the interesting bit is that the two-pass")

print(" max/sum reduction never left registers.")

行列乗算カーネルを拡張し、バイアス加算と GELU 活性化関数をレジスタに保持されたアキュムレータに直接融合させます。これにより中間的なグローバルメモリへのアクセスを削減し、最終出力テンソルを書き出す前にエピローグ処理を完了させることで、統合実装の性能を eager な PyTorch 実行と比較評価します。

また、フラグメントレベルでの最大値および和の計算を用いた行ごとのソフトマックスカーネルを実装しました。これにより正規化プロセスの大部分をレジスタ内で完結させ、メモリアクセスのオーバーヘッドを抑えています。

@tilelang.jit(out_idx=[-1])

def make_flash_attn(batch: int, heads: int, seq_len: int, dim: int,

is_causal: bool = False,

block_M: int = 64, block_N: int = 64,

num_stages: int = 1, threads: int = 128,

dtype: str = "float16", accum_dtype: str = "float"):

scale = 1.0 / math.sqrt(dim)

shape = [batch, seq_len, heads, dim]

NEG = -1.0e30

@T.prim_func

def main(Q: T.Tensor(shape, dtype),

K: T.Tensor(shape, dtype),

V: T.Tensor(shape, dtype),

O: T.Tensor(shape, dtype)):

with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch,

threads=threads) as (bx, by, bz):

Q_shared = T.alloc_shared([block_M, dim], dtype)

K_shared = T.alloc_shared([block_N, dim], dtype)

V_shared = T.alloc_shared([block_N, dim], dtype)

acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)

acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)

acc_o = T.alloc_fragment([block_M, dim], accum_dtype)

m_prev = T.alloc_fragment([block_M], accum_dtype)

m_cur = T.alloc_fragment([block_M], accum_dtype)

alpha = T.alloc_fragment([block_M], accum_dtype)

p_sum = T.alloc_fragment([block_M], accum_dtype)

logsum = T.alloc_fragment([block_M], accum_dtype)

T.copy(Q[bz, bx * block_M:(bx + 1) * block_M, by, :], Q_shared)

T.fill(acc_o, 0)

T.fill(logsum, 0)

T.fill(m_cur, NEG)

loop_range = (T.ceildiv((bx + 1) * block_M, block_N)

if is_causal else T.ceildiv(seq_len, block_N))

for k in T.Pipelined(loop_range, num_stages=num_stages):

T.copy(K[bz, k * block_N:(k + 1) * block_N, by, :], K_shared)

if is_causal:

for i, j in T.Parallel(block_M, block_N):

acc_s[i, j] = T.if_then_else(

bx * block_M + i >= k * block_N + j, 0.0, NEG)

else:

T.clear(acc_s)

T.gemm(Q_shared, K_shared, acc_s, transpose_B=True)

T.copy(V[bz, k * block_N:(k + 1) * block_N, by, :], V_shared)

T.copy(m_cur, m_prev)

T.reduce_max(acc_s, m_cur, dim=1, clear=False)

for i in T.Parallel(block_M):

alpha[i] = T.exp((m_prev[i] - m_cur[i]) * scale)

for i, j in T.Parallel(block_M, block_N):

acc_s[i, j] = T.exp((acc_s[i, j] - m_cur[i]) * scale)

T.reduce_sum(acc_s, p_sum, dim=1)

for i in T.Parallel(block_M):

logsum[i] = logsum[i] * alpha[i] + p_sum[i]

for i, j in T.Parallel(block_M, dim):

acc_o[i, j] = acc_o[i, j] * alpha[i]

T.copy(acc_s, acc_s_cast)

T.gemm(acc_s_cast, V_shared, acc_o)

for i, j in T.Parallel(block_M, dim):

acc_o[i, j] = acc_o[i, j] / logsum[i]

T.copy(acc_o, O[bz, bx * block_M:(bx + 1) * block_M, by, :])

return main

def section_6():

banner("6. FLASHATTENTION — fused attention forward")

B, H, S, D = 4, 8, 1024, 64

q = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16)

k = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16)

v = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16)

def torch_ref(causal: bool):

qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v))

o = F.scaled_dot_product_attention(qt, kt, vt, is_causal=causal)

return o.transpose(1, 2).contiguous()

for causal in (False, True):

tag = "causal" if causal else "full"

attn = make_flash_attn(B, H, S, D, is_causal=causal,

block_M=64, block_N=64,

num_stages=1 if SM < 80 else 2, threads=128)

o = attn(q, k, v)

ref = torch_ref(causal)

check(o, ref, f"flash_attn ({tag})", tol=3e-2)

flops = 4 * B * H * S * S * D * (0.5 if causal else 1.0)

ms = bench(lambda: attn(q, k, v), warmup=5, rep=20)

ms_ref = bench(lambda: torch_ref(causal), warmup=5, rep=20)

print(f" {tag:>6}: tilelang {ms:6.3f} ms "

f"({flops/(ms*1e-3)/1e12:5.2f} TFLOP/s) "

f"torch SDPA {ms_ref:6.3f} ms "

f"({flops/(ms_ref*1e-3)/1e12:5.2f} TFLOP/s)")

print(" ~70 lines of Python for a fused, causal, tensor-core attention.")

print(" The upstream repo pushes the same structure to FlashMLA-level perf")

print(" on H100 with warp specialisation and TMA.")

私たちは、融合された FlashAttention のカーネルを実装しました。