Designing High-Performance GPU Kernels with TileLang: Tensor-Core GEMM, Fused Softmax, FlashAttention, and Autotuning

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