A Coding Guide to NVIDIA’s Tile-Based GPU Programming: From cuTile and Triton Kernels to Flash Attention
In this tutorial, we explore TileGym GPU programming by building a practical Colab workflow that runs across different hardware conditions. We begin by probing the available CUDA environment, checking whether NVIDIA cuTile runs directly, and falling back to Triton when standard Colab GPUs lack the required cuTile stack. Through this setup, we learn the core tile-programming idea: instead of writing code for one thread at a time, we operate on entire data tiles, load them into the kernel, compute on them efficiently, and store the results back. We use this model to implement vector addition, fused GELU, row-wise softmax, tiled matrix multiplication, and flash attention, while comparing each result against PyTorch for correctness and benchmarking. CUDA Environment Probe Copy CodeCopiedUse a different Browser import os, sys, math, time, textwrap def rule(t=””): print(“n” + “=” * 78) if t: print(t) print(“=” * 78) rule(“0. ENVIRONMENT PROBE”) try: import torch except ImportError: print(“Installing torch …”) os.system(f”{sys.executable} -m pip install -q torch”) import torch HAS_CUDA = torch.cuda.is_available() DEV = “cuda” if HAS_CUDA else “cpu” cc = (0, 0) if HAS_CUDA: cc = torch.cuda.get_device_capability() print(f”GPU : {torch.cuda.get_device_name(0)}”) print(f”Compute capability : {cc[0]}.{cc[1]}”) print(f”Torch CUDA runtime : {torch.version.cuda}”) print(f”Driver / torch : {torch.__version__}”) else: print(“No CUDA GPU found. In Colab: Runtime -> Change runtime type -> GPU (T4).”) print(“The tutorial will still run its correctness math on CPU where possible.”) CUTILE_HW_OK = HAS_CUDA and cc[0] >= 8 CUDA_MAJOR = int((torch.version.cuda or “0”).split(“.”)[0]) if HAS_CUDA else 0 CUTILE_TOOLKIT_OK = CUDA_MAJOR >= 13 rule(“1. ATTEMPTING REAL cuTile (NVIDIA CUDA Tile) BACKEND”) ct = None CUTILE_READY = False if CUTILE_HW_OK and CUTILE_TOOLKIT_OK: try: import cuda.tile as ct CUTILE_READY = True print(“cuda.tile is already importable.”) except Exception: print(“Installing cuda-tile[tileiras] (this can take a while)…”) os.system(f”{sys.executable} -m pip install -q ‘cuda-tile[tileiras]’ cupy-cuda13x”) try: import cuda.tile as ct CUTILE_READY = True except Exception as e: print(“cuTile import still failed:”, repr(e)) else: reasons = [] if not HAS_CUDA: reasons.append(“no CUDA GPU”) if HAS_CUDA and cc[0] < 8: reasons.append(f”compute capability {cc[0]}.{cc[1]} < 8.0 (Turing/T4 unsupported)”) if not CUTILE_TOOLKIT_OK: reasons.append(f”CUDA {torch.version.cuda} < 13.1 required by tileiras”) print(“Skipping real cuTile install because:”, “; “.join(reasons) + “.”) print(“This is expected on a standard Colab T4 — we fall back to Triton below,”) print(“which teaches the exact same tile-based programming model.”) if CUTILE_READY: BACKEND = “cutile” else: try: import triton, triton.language as tl BACKEND = “triton” if HAS_CUDA else “torch” except ImportError: if HAS_CUDA: print(“Installing triton …”) os.system(f”{sys.executable} -m pip install -q triton”) try: import triton, triton.language as tl BACKEND = “triton” except Exception: BACKEND = “torch” else: BACKEND = “torch” rule(f”ACTIVE EXECUTION BACKEND: {BACKEND.upper()}”) print({ “cutile”: “Running NVIDIA cuTile kernels on your Ampere+/CUDA13 GPU. Nice hardware!”, “triton”: “Running Triton tile kernels on your GPU (the standard Colab path).”, “torch”: “No usable GPU kernel backend; showing reference math on CPU only.”, }[BACKEND]) print(textwrap.dedent(“”” —————————————————————— SIMT (classic CUDA) | TILE model (cuTile / Triton) —————————————————————— You write code for ONE | You write code for ONE BLOCK that thread. You compute a global | owns a whole TILE (e.g. 1024 elems index, bounds-check it, and | or a 128×128 sub-matrix). You load touch a single element. | the tile, do math on the WHOLE tile, | store it. The compiler maps the tile C[i] = A[i] + B[i] | onto threads / tensor cores for you. —————————————————————— cuTile primitives: ct.bid(0), ct.load(…), ct.store(…), a @ b, ct.launch Triton primitives: tl.program_id, tl.load, tl.store, tl.dot, grid[…] Same idea, two spellings. Below, every kernel is shown in BOTH. “””)) CUTILE_SOURCE = { “vector_add”: ”’ import cuda.tile as ct @ct.kernel def vector_add(a, b, c, tile_size: ct.Constant[int]): pid = ct.bid(0) a_tile = ct.load(a, index=(pid,), shape=(tile_size,)) b_tile = ct.load(b, index=(pid,), shape=(tile_size,)) ct.store(c, index=(pid,), tile=a_tile + b_tile) ”’, “matmul”: ”’ import cuda.tile as ct @ct.kernel def matmul(A, B, C, K: ct.Constant[int], BM: ct.Constant[int], BN: ct.Constant[int], BK: ct.Constant[int]): m, n = ct.bid(0), ct.bid(1) acc = ct.zeros((BM, BN), dtype=ct.float32) for k in range(ct.cdiv(K, BK)): a = ct.load(A, index=(m, k), shape=(BM, BK)) b = ct.load(B, index=(k, n), shape=(BK, BN)) acc = a @ b + acc ct.store(C, index=(m, n), tile=acc) ”’} We begin by setting up the environment, importing the required libraries, and checking whether CUDA is available on the current runtime. We inspect the GPU capabilities, CUDA version, and PyTorch setup to determine whether the real cuTile backend is usable. We then select the active execution backend, explain the tile programming model, and store reference cuTile kernel source strings for comparison. Defining Triton Kernels Copy CodeCopiedUse a different Browser if BACKEND == “triton”: @triton.jit def _vadd_kernel(a_ptr, b_ptr, c_ptr, n, BLOCK: tl.constexpr): pid = tl.program_id(0) offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n a = tl.load(a_ptr + offs, mask=mask) b = tl.load(b_ptr + offs, mask=mask) tl.store(c_ptr + offs, a + b, mask=mask) @triton.jit def _fused_gelu_kernel(x_ptr, w_ptr, b_ptr, o_ptr, n, BLOCK: tl.constexpr): pid = tl.program_id(0) offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n x = tl.load(x_ptr + offs, mask=mask) w = tl.load(w_ptr + offs, mask=mask) b = tl.load(b_ptr + offs, mask=mask) h = x * w + b c = 0.7978845608028654 z = c * (h + 0.044715 * h * h * h) e = tl.exp(-2.0 * z) tanh = (1.0 – e) / (1.0 + e) g = 0.5 * h * (1.0 + tanh) tl.store(o_ptr + offs, g, mask=mask) @triton.jit def _softmax_kernel(x_ptr, o_ptr, stride, n_cols, BLOCK: tl.constexpr): row = tl.program_id(0) cols = tl.arange(0, BLOCK) mask = cols < n_cols ptr = x_ptr + row * stride + cols x = tl.load(ptr, mask=mask, other=-float(“inf”)) x = x – tl.max(x, axis=0) num = tl.exp(x) den = tl.sum(num, axis=0) tl.store(o_ptr + row * stride + cols, num / den, mask=mask) @triton.jit def _matmul_kernel(A, B, C, M, N, K, sam, sak, sbk, sbn, scm, scn, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr): pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BM + tl.arange(0, BM) offs_n = pid_n * BN + tl.arange(0, BN) offs_k = tl.arange(0, BK) a_ptr = A + offs_m[:, None] * sam + offs_k[None, :] * sak b_ptr

