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

In this tutorial, we explore TileLang as a high-level Python domain-specific language for designing and compiling performance-oriented GPU kernels through TVM. We begin by validating the CUDA environment and establishing reusable benchmarking and numerical-verification utilities, then progressively implement vector addition, tiled tensor-core matrix multiplication, schedule exploration, fused GEMM epilogues, row-wise softmax, and FlashAttention. Throughout the tutorial, we work directly with TileLang’s shared-memory tiles, register fragments, pipelined loops, parallel iteration primitives, reductions, and tensor-core GEMM operators while allowing the compiler to manage thread mapping, memory layouts, synchronization, vectorization, and low-level CUDA instruction generation. We also compare our kernels against PyTorch and cuBLAS baselines, inspect generated CUDA source, evaluate memory and compute throughput, and use autotuning to identify architecture-dependent kernel configurations.

Copy CodeCopiedUse a different Browser

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():

"""Install tilelang if missing. Stable wheel first, nightly as a fallback."""

try:

import tilelang

return

except ImportError:

pass

print(">> installing tilelang (this pulls a bundled TVM, ~1-3 min)\n")

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

try:

import tilelang

return

except ImportError:

print(">> stable wheel unusable, trying nightly channel")

_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:

"""Median-ish latency in milliseconds, measured with CUDA events."""

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:

"""Relative-Frobenius-norm check. Far more meaningful than atol for fp16."""

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(), "No GPU. Runtime -> Change runtime type -> 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} SMs, "

f"{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 caches compiled kernels in ~/.tilelang/cache, so a")

print(" second run of this cell is dramatically faster.")

@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 — vector add + generated 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 <- both are pure bandwidth, so a tie"

f" is the correct outcome")

src = add.get_kernel_source()

print("\n --- generated device code (first 32 lines) ---")

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

print(" " + line)

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

We configure the Google Colab CUDA environment, install TileLang with a nightly fallback, and import the required PyTorch and TileLang modules. We define reusable benchmarking, validation, and reporting utilities that measure kernel latency and compare numerical outputs using relative error. We then implement a TileLang vector-add kernel, execute it on the GPU, compare its bandwidth with PyTorch, and inspect the CUDA source generated by the compiler.

Copy CodeCopiedUse a different Browser

@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. MEMORY HIERARCHY — tiled 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. KNOBS — sweeping the schedule by hand")

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.")

We implement a tiled tensor-core matrix multiplication kernel that moves input tiles through global memory, shared memory, and register fragments. We control tile dimensions, pipeline stages, thread counts, and L2 swizzling while allowing TileLang to generate tensor-core instructions, synchronization, and memory-transfer logic. We then benchmark several schedule configurations, verify their numerical accuracy, and identify the highest-performing architecture-dependent kernel configuration.

Copy CodeCopiedUse a different Browser

@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.")

We extend the matrix multiplication kernel by fusing bias addition and the GELU activation directly into the register-resident accumulator. We reduce intermediate global-memory traffic by completing the epilogue before writing the final output tensor and compare the fused implementation with eager PyTorch execution. We also implement a row-wise softmax kernel using fragment-level maximum and sum reductions while keeping the normalization process largely within registers.

Copy CodeCopiedUse a different Browser

@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.")

We implement a fused FlashAtte