Coding’s Guide to NVIDIA’s Tile-Based GPU Programming: From cuTile and Triton Kernels to 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, 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 = B + offs_k[:, None] * sbk + offs_n[None, :] * sbn
       acc = tl.zeros((BM, BN), dtype=tl.float32)
       for k in range(0, K, BK):
           a = tl.load(a_ptr, mask=offs_k[None, :] < K - k, other=0.0)
           b = tl.load(b_ptr, mask=offs_k[:, None] < K - k, other=0.0)
           acc += tl.dot(a, b)
           a_ptr += BK * sak
           b_ptr += BK * sbk
       c_ptr = C + offs_m[:, None] * scm + offs_n[None, :] * scn
       cmask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
       tl.store(c_ptr, acc.to(C.dtype.element_ty), mask=cmask)
   @triton.jit
   def _flash_kernel(Q, K, V, O, sqz, skz, svz, soz,
                     L, D, scale,
                     BL: tl.constexpr, BD: tl.constexpr):
       pid_l = tl.program_id(0)
       z     = tl.program_id(1)
       offs_l = pid_l * BL + tl.arange(0, BL)
       offs_d = tl.arange(0, BD)
       q_ptr = Q + z * sqz + offs_l[:, None] * D + offs_d[None, :]
       q = tl.load(q_ptr, mask=offs_l[:, None] < L, other=0.0)
       m_i = tl.full((BL,), -float("inf"), dtype=tl.float32)
       l_i = tl.zeros((BL,), dtype=tl.float32)
       acc = tl.zeros((BL, BD), dtype=tl.float32)
       for start in range(0, L, BL):
           offs_k = start + tl.arange(0, BL)
           k_ptr = K + z * skz + offs_k[:, None] * D + offs_d[None, :]
           v_ptr = V + z * svz + offs_k[:, None] * D + offs_d[None, :]
           k = tl.load(k_ptr, mask=offs_k[:, None] < L, other=0.0)
           v = tl.load(v_ptr, mask=offs_k[:, None] < L, other=0.0)
           s = tl.dot(q, tl.trans(k)) * scale
           s = tl.where(offs_k[None, :] < L, s, -float("inf"))
           m_ij = tl.maximum(m_i, tl.max(s, axis=1))
           p    = tl.exp(s - m_ij[:, None])
           alpha = tl.exp(m_i - m_ij)
           l_i  = l_i * alpha + tl.sum(p, axis=1)
           acc  = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
           m_i  = m_ij
       acc = acc / l_i[:, None]
       o_ptr = O + z * soz + offs_l[:, None] * D + offs_d[None, :]
       tl.store(o_ptr, acc.to(O.dtype.element_ty), mask=offs_l[:, None] < L)
   def run_vadd(a, b):
       c = torch.empty_like(a); n = a.numel()
       grid = (triton.cdiv(n, 1024),)
       _vadd_kernel[grid](a, b, c, n, BLOCK=1024)
       return c
   def run_fused_gelu(x, w, b):
       o = torch.empty_like(x); n = x.numel()
       grid = (triton.cdiv(n, 1024),)
       _fused_gelu_kernel[grid](x, w, b, o, n, BLOCK=1024)
       return o
   def run_softmax(x):
       m, ncols = x.shape
       o = torch.empty_like(x)
       BLOCK = triton.next_power_of_2(ncols)
       _softmax_kernel[(m,)](x, o, x.stride(0), ncols, BLOCK=BLOCK)
       return o
   def run_matmul(a, b):
       M, K = a.shape; K2, N = b.shape
       c = torch.empty((M, N), device=a.device, dtype=a.dtype)
       BM = BN = 64; BK = 32
       grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
       _matmul_kernel[grid](a, b, c, M, N, K,
                            a.stride(0), a.stride(1), b.stride(0), b.stride(1),
                            c.stride(0), c.stride(1), BM=BM, BN=BN, BK=BK)
       return c
   def run_flash(q, k, v):
       Z, L, D = q.shape
       o = torch.empty_like(q)
       scale = 1.0 / math.sqrt(D)
       BL = 64
       grid = (triton.cdiv(L, BL), Z)
       _flash_kernel[grid](q, k, v, o,
                           q.stride(0), k.stride(0), v.stride(0), o.stride(0),
                           L, D, scale, BL=BL, BD=D)
       return o
else:
   def run_vadd(a, b):            return a + b
   def run_fused_gelu(x, w, b):   return torch.nn.functional.gelu(x * w + b, approximate="tanh")
   def run_softmax(x):            return torch.softmax(x, dim=-1)
   def run_matmul(a, b):          return a @ b
   def run_flash(q, k, v):        return torch.nn.functional.scaled_dot_product_attention(q, k, v)