TileGym is a Colab workflow designed to teach GPU programming by comparing two methods: NVIDIA’s cuTile and the Triton compiler. The tutorial checks the local environment for CUDA availability and hardware capability, switching to Triton if the required cuTile stack is missing. It demonstrates tile-based programming by processing entire data blocks rather than individual threads. The examples cover vector addition, fused GELU, row-wise softmax, tiled matrix multiplication, and flash attention. Each result is verified against PyTorch and benchmarked.
Checking the Hardware
The script begins by importing libraries and verifying if CUDA is active. It inspects the GPU name, compute capability, and the PyTorch version. If the hardware lacks the necessary features, the system falls back to CPU calculations.
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 128x128 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)
'''}
The code attempts to import the `cuda.tile` module. If the hardware is an Ampere-class GPU running CUDA 13 or higher, it proceeds with the native backend. If the environment is standard Colab T4 hardware, the script skips the installation and switches to Triton. The output explains the difference between the SIMT model, where code targets a single thread, and the tile model, where code targets a block of data.
Defining Triton Kernels
When the native backend is unavailable, the tutorial defines kernels using the Triton compiler. These functions handle vector addition, fused GELU activation, softmax, matrix multiplication, and flash attention.
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, strideSource Read original →Related articles
