{"slug": "designing-high-performance-gpu-kernels-with-tilelang-tensor-core-gemm-fused-and", "title": "Designing High-Performance GPU Kernels with TileLang: Tensor-Core GEMM, Fused Softmax, FlashAttention, and Autotuning", "summary": "TileLang, a high-level Python domain-specific language for designing GPU kernels through TVM, enables developers to implement tensor-core GEMM, fused softmax, FlashAttention, and autotuning while managing shared-memory tiles, register fragments, and pipelined loops. The tutorial validates CUDA environments, benchmarks kernels against PyTorch and cuBLAS baselines, and uses autotuning to identify architecture-dependent configurations.", "body_md": "In this tutorial, we explore[ TileLang](https://github.com/tile-ai/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.\n\n``` python\nimport os\nimport sys\nimport math\nimport time\nimport subprocess\nimport traceback\ndef _sh(cmd: str) -> int:\n   print(f\"$ {cmd}\", flush=True)\n   return subprocess.run(cmd, shell=True).returncode\ndef _bootstrap():\n   \"\"\"Install tilelang if missing. Stable wheel first, nightly as a fallback.\"\"\"\n   try:\n       import tilelang\n       return\n   except ImportError:\n       pass\n   print(\">> installing tilelang (this pulls a bundled TVM, ~1-3 min)\\n\")\n   _sh(f\"{sys.executable} -m pip install -q tilelang\")\n   try:\n       import tilelang\n       return\n   except ImportError:\n       print(\">> stable wheel unusable, trying nightly channel\")\n       _sh(f\"{sys.executable} -m pip install -q tilelang \"\n           f\"-f https://tile-ai.github.io/whl/nightly\")\n       import tilelang\nif os.path.isdir(\"/usr/local/cuda\"):\n   os.environ.setdefault(\"CUDA_HOME\", \"/usr/local/cuda\")\n   os.environ[\"PATH\"] = os.environ.get(\"PATH\", \"\") + \":/usr/local/cuda/bin\"\n_bootstrap()\nimport torch\nimport torch.nn.functional as F\nimport tilelang\nimport tilelang.language as T\ndef banner(title: str):\n   print(\"\\n\" + \"=\" * 78)\n   print(f\"  {title}\")\n   print(\"=\" * 78, flush=True)\ndef bench(fn, warmup: int = 10, rep: int = 50) -> float:\n   \"\"\"Median-ish latency in milliseconds, measured with CUDA events.\"\"\"\n   for _ in range(warmup):\n       fn()\n   torch.cuda.synchronize()\n   start, end = torch.cuda.Event(True), torch.cuda.Event(True)\n   start.record()\n   for _ in range(rep):\n       fn()\n   end.record()\n   torch.cuda.synchronize()\n   return start.elapsed_time(end) / rep\ndef check(actual: torch.Tensor, ref: torch.Tensor, name: str, tol: float = 2e-2) -> bool:\n   \"\"\"Relative-Frobenius-norm check. Far more meaningful than atol for fp16.\"\"\"\n   a, r = actual.float(), ref.float()\n   rel = (a - r).norm() / r.norm().clamp_min(1e-12)\n   amax = (a - r).abs().max().item()\n   ok = bool(rel < tol) and math.isfinite(amax)\n   print(f\"   [{'PASS' if ok else 'FAIL'}] {name}: rel_err={rel:.3e}  max_abs={amax:.3e}\")\n   return ok\nbanner(\"0. ENVIRONMENT\")\nassert torch.cuda.is_available(), \"No GPU. Runtime -> Change runtime type -> GPU.\"\nDEV = torch.device(\"cuda\")\nPROPS = torch.cuda.get_device_properties(0)\nCC = torch.cuda.get_device_capability(0)\nSM = CC[0] * 10 + CC[1]\nprint(f\"   tilelang   : {getattr(tilelang, '__version__', 'unknown')}\")\nprint(f\"   torch      : {torch.__version__}\")\nprint(f\"   GPU        : {PROPS.name}  (sm_{SM}, {PROPS.multi_processor_count} SMs, \"\n     f\"{PROPS.total_memory/2**30:.1f} GiB)\")\nSMEM_CAP = 48 * 1024 if SM < 80 else 96 * 1024\nDEFAULT_STAGES = 2 if SM < 80 else 3\nprint(f\"   smem budget: {SMEM_CAP//1024} KB/block, default num_stages={DEFAULT_STAGES}\")\nprint(\"   note: tilelang caches compiled kernels in ~/.tilelang/cache, so a\")\nprint(\"         second run of this cell is dramatically faster.\")\n@tilelang.jit(out_idx=[-1])\ndef make_vector_add(N: int, block_N: int = 256, dtype: str = \"float32\"):\n   @T.prim_func\n   def main(A: T.Tensor((N,), dtype),\n            B: T.Tensor((N,), dtype),\n            C: T.Tensor((N,), dtype)):\n       with T.Kernel(T.ceildiv(N, block_N), threads=256) as bx:\n           for i in T.Parallel(block_N):\n               C[bx * block_N + i] = A[bx * block_N + i] + B[bx * block_N + i]\n   return main\ndef section_1():\n   banner(\"1. HELLO, TILE — vector add + generated CUDA\")\n   N = 1 << 22\n   add = make_vector_add(N, block_N=256)\n   a = torch.randn(N, device=DEV)\n   b = torch.randn(N, device=DEV)\n   c = add(a, b)\n   check(c, a + b, \"vector_add\")\n   ms = bench(lambda: add(a, b))\n   gbs = 3 * N * 4 / (ms * 1e-3) / 1e9\n   ref_ms = bench(lambda: a + b)\n   print(f\"   tilelang: {ms*1e3:8.1f} us  ({gbs:6.1f} GB/s)\")\n   print(f\"   torch   : {ref_ms*1e3:8.1f} us   <- both are pure bandwidth, so a tie\"\n         f\" is the correct outcome\")\n   src = add.get_kernel_source()\n   print(\"\\n   --- generated device code (first 32 lines) ---\")\n   for line in src.splitlines()[:32]:\n       print(\"   \" + line)\n   print(\"   ---------------------------------------------\")\n```\n\nWe 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.\n\n``` python\n@tilelang.jit(out_idx=[-1])\ndef make_matmul(M: int, N: int, K: int,\n               block_M: int = 128, block_N: int = 128, block_K: int = 32,\n               num_stages: int = 3, threads: int = 128,\n               use_swizzle: bool = False,\n               dtype: str = \"float16\", accum_dtype: str = \"float\"):\n   @T.prim_func\n   def main(A: T.Tensor((M, K), dtype),\n            B: T.Tensor((K, N), dtype),\n            C: T.Tensor((M, N), dtype)):\n       with T.Kernel(T.ceildiv(N, block_N),\n                     T.ceildiv(M, block_M),\n                     threads=threads) as (bx, by):\n           A_shared = T.alloc_shared((block_M, block_K), dtype)\n           B_shared = T.alloc_shared((block_K, block_N), dtype)\n           C_local = T.alloc_fragment((block_M, block_N), accum_dtype)\n           if use_swizzle:\n               T.use_swizzle(panel_size=10, enable=True)\n           T.clear(C_local)\n           for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):\n               T.copy(A[by * block_M, ko * block_K], A_shared)\n               T.copy(B[ko * block_K, bx * block_N], B_shared)\n               T.gemm(A_shared, B_shared, C_local)\n           T.copy(C_local, C[by * block_M, bx * block_N])\n   return main\ndef smem_bytes(block_M, block_N, block_K, stages, itemsize=2):\n   return (block_M * block_K + block_K * block_N) * itemsize * stages\ndef section_2():\n   banner(\"2. MEMORY HIERARCHY — tiled tensor-core GEMM\")\n   M = N = K = 2048\n   bm, bn, bk, st = 128, 128, 32, DEFAULT_STAGES\n   while smem_bytes(bm, bn, bk, st) > SMEM_CAP and st > 1:\n       st -= 1\n   print(f\"   config: {bm}x{bn}x{bk}, num_stages={st}, \"\n         f\"smem={smem_bytes(bm,bn,bk,st)/1024:.0f} KB\")\n   kernel = make_matmul(M, N, K, bm, bn, bk, num_stages=st, threads=128)\n   a = torch.randn(M, K, device=DEV, dtype=torch.float16)\n   b = torch.randn(K, N, device=DEV, dtype=torch.float16)\n   c = kernel(a, b)\n   check(c, a @ b, \"matmul 2048^3\")\n   flops = 2 * M * N * K\n   ms = bench(lambda: kernel(a, b))\n   ms_ref = bench(lambda: a @ b)\n   print(f\"   tilelang : {ms:7.3f} ms  ->  {flops/(ms*1e-3)/1e12:6.2f} TFLOP/s\")\n   print(f\"   cuBLAS   : {ms_ref:7.3f} ms  ->  {flops/(ms_ref*1e-3)/1e12:6.2f} TFLOP/s\")\n   print(f\"   ratio    : {ms_ref/ms*100:5.1f}% of cuBLAS from ~20 lines of Python\")\n   src = kernel.get_kernel_source()\n   for needle in (\"mma.sync\", \"wgmma\", \"ldmatrix\", \"cp.async\", \"tl::gemm\"):\n       if needle in src:\n           print(f\"   emitted: {needle}\")\n   return kernel\ndef section_3():\n   banner(\"3. KNOBS — sweeping the schedule by hand\")\n   M = N = K = 2048\n   a = torch.randn(M, K, device=DEV, dtype=torch.float16)\n   b = torch.randn(K, N, device=DEV, dtype=torch.float16)\n   flops = 2 * M * N * K\n   candidates = [\n       (64,  64,  32, 2, 128, False),\n       (128, 128, 32, 2, 128, False),\n       (128, 128, 32, 2, 128, True),\n       (128, 128, 32, 3, 128, False),\n       (128, 128, 64, 2, 256, False),\n       (128, 256, 32, 2, 256, False),\n   ]\n   print(f\"   {'tile':>16} {'stg':>4} {'thr':>4} {'swz':>4} {'smem':>7} \"\n         f\"{'ms':>8} {'TFLOP/s':>9}\")\n   results = []\n   for bm, bn, bk, st, thr, swz in candidates:\n       need = smem_bytes(bm, bn, bk, st)\n       tag = f\"{bm}x{bn}x{bk}\"\n       if need > SMEM_CAP:\n           print(f\"   {tag:>16} {st:>4} {thr:>4} {str(swz):>4} {need//1024:>5}KB \"\n                 f\"  SKIPPED (over smem budget)\")\n           continue\n       try:\n           k = make_matmul(M, N, K, bm, bn, bk, num_stages=st,\n                           threads=thr, use_swizzle=swz)\n           c = k(a, b)\n           ok = (c - (a @ b)).float().norm() / (a @ b).float().norm() < 2e-2\n           ms = bench(lambda: k(a, b), warmup=5, rep=20)\n           results.append((ms, tag, st, thr, swz))\n           print(f\"   {tag:>16} {st:>4} {thr:>4} {str(swz):>4} {need//1024:>5}KB \"\n                 f\"{ms:>8.3f} {flops/(ms*1e-3)/1e12:>9.2f}\"\n                 f\"{'' if ok else '   <- NUMERICALLY WRONG'}\")\n       except Exception as e:\n           print(f\"   {tag:>16} {st:>4} {thr:>4} {str(swz):>4}     -  \"\n                 f\"failed: {type(e).__name__}\")\n   if results:\n       best = min(results)\n       print(f\"\\n   winner: {best[1]}, stages={best[2]}, threads={best[3]}, \"\n             f\"swizzle={best[4]}  ({best[0]:.3f} ms)\")\n   print(\"   Takeaway: the best schedule is arch- and shape-dependent, which is\")\n   print(\"   exactly why section 7 exists.\")\n```\n\nWe 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.\n\n``` python\n@tilelang.jit(out_idx=[-1])\ndef make_matmul_bias_gelu(M: int, N: int, K: int,\n                         block_M: int = 128, block_N: int = 128, block_K: int = 32,\n                         num_stages: int = 2, threads: int = 128,\n                         dtype: str = \"float16\", accum_dtype: str = \"float\"):\n   @T.prim_func\n   def main(A: T.Tensor((M, K), dtype),\n            B: T.Tensor((K, N), dtype),\n            Bias: T.Tensor((N,), dtype),\n            C: T.Tensor((M, N), dtype)):\n       with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M),\n                     threads=threads) as (bx, by):\n           A_shared = T.alloc_shared((block_M, block_K), dtype)\n           B_shared = T.alloc_shared((block_K, block_N), dtype)\n           Bias_shared = T.alloc_shared((block_N,), dtype)\n           C_local = T.alloc_fragment((block_M, block_N), accum_dtype)\n           T.clear(C_local)\n           T.copy(Bias[bx * block_N], Bias_shared)\n           for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):\n               T.copy(A[by * block_M, ko * block_K], A_shared)\n               T.copy(B[ko * block_K, bx * block_N], B_shared)\n               T.gemm(A_shared, B_shared, C_local)\n           for i, j in T.Parallel(block_M, block_N):\n               C_local[i, j] = C_local[i, j] + Bias_shared[j]\n           for i, j in T.Parallel(block_M, block_N):\n               C_local[i, j] = C_local[i, j] / (\n                   1.0 + T.exp(-1.5957691216 * (\n                       C_local[i, j] + 0.044715 * C_local[i, j]\n                       * C_local[i, j] * C_local[i, j])))\n           T.copy(C_local, C[by * block_M, bx * block_N])\n   return main\ndef section_4():\n   banner(\"4. EPILOGUE FUSION — GEMM + bias + GELU in one kernel\")\n   M, N, K = 4096, 4096, 1024\n   st = DEFAULT_STAGES\n   while smem_bytes(128, 128, 32, st) > SMEM_CAP and st > 1:\n       st -= 1\n   fused = make_matmul_bias_gelu(M, N, K, 128, 128, 32, num_stages=st, threads=128)\n   a = (torch.randn(M, K, device=DEV, dtype=torch.float16) / K ** 0.25)\n   b = (torch.randn(K, N, device=DEV, dtype=torch.float16) / K ** 0.25)\n   bias = torch.randn(N, device=DEV, dtype=torch.float16)\n   out = fused(a, b, bias)\n   ref = F.gelu(((a @ b).float() + bias.float()), approximate=\"tanh\")\n   check(out, ref, \"matmul+bias+gelu\", tol=3e-2)\n   ms_f = bench(lambda: fused(a, b, bias))\n   ms_e = bench(lambda: F.gelu(a @ b + bias, approximate=\"tanh\"))\n   print(f\"   fused (1 kernel)   : {ms_f:7.3f} ms\")\n   print(f\"   torch (3 kernels)  : {ms_e:7.3f} ms\")\n   print(f\"   speedup            : {ms_e/ms_f:5.2f}x\")\n   print(f\"   HBM traffic saved  : ~{2*M*N*2/2**20:.0f} MiB of intermediate reads/writes\")\n@tilelang.jit(out_idx=[-1])\ndef make_softmax(M: int, N: int, block_M: int = 4, threads: int = 128,\n                dtype: str = \"float16\", accum_dtype: str = \"float\"):\n   @T.prim_func\n   def main(X: T.Tensor((M, N), dtype),\n            Y: T.Tensor((M, N), dtype)):\n       with T.Kernel(T.ceildiv(M, block_M), threads=threads) as bx:\n           X_shared = T.alloc_shared((block_M, N), dtype)\n           X_local = T.alloc_fragment((block_M, N), accum_dtype)\n           row_max = T.alloc_fragment((block_M,), accum_dtype)\n           row_sum = T.alloc_fragment((block_M,), accum_dtype)\n           T.copy(X[bx * block_M, 0], X_shared)\n           T.copy(X_shared, X_local)\n           T.reduce_max(X_local, row_max, dim=1, clear=True)\n           for i, j in T.Parallel(block_M, N):\n               X_local[i, j] = T.exp(X_local[i, j] - row_max[i])\n           T.reduce_sum(X_local, row_sum, dim=1)\n           for i, j in T.Parallel(block_M, N):\n               X_local[i, j] = X_local[i, j] / row_sum[i]\n           T.copy(X_local, Y[bx * block_M, 0])\n   return main\ndef section_5():\n   banner(\"5. REDUCTIONS — row-wise softmax\")\n   M, N = 8192, 1024\n   sm = make_softmax(M, N, block_M=4, threads=128)\n   x = torch.randn(M, N, device=DEV, dtype=torch.float16) * 3.0\n   y = sm(x)\n   check(y, torch.softmax(x.float(), dim=-1), \"softmax\")\n   ms = bench(lambda: sm(x))\n   ms_ref = bench(lambda: torch.softmax(x, dim=-1))\n   gbs = 2 * M * N * 2 / (ms * 1e-3) / 1e9\n   print(f\"   tilelang: {ms*1e3:7.1f} us  ({gbs:6.1f} GB/s)\")\n   print(f\"   torch   : {ms_ref*1e3:7.1f} us\")\n   print(\"   Both are memory bound; the interesting bit is that the two-pass\")\n   print(\"   max/sum reduction never left registers.\")\n```\n\nWe 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.\n\n``` python\n@tilelang.jit(out_idx=[-1])\ndef make_flash_attn(batch: int, heads: int, seq_len: int, dim: int,\n                   is_causal: bool = False,\n                   block_M: int = 64, block_N: int = 64,\n                   num_stages: int = 1, threads: int = 128,\n                   dtype: str = \"float16\", accum_dtype: str = \"float\"):\n   scale = 1.0 / math.sqrt(dim)\n   shape = [batch, seq_len, heads, dim]\n   NEG = -1.0e30\n   @T.prim_func\n   def main(Q: T.Tensor(shape, dtype),\n            K: T.Tensor(shape, dtype),\n            V: T.Tensor(shape, dtype),\n            O: T.Tensor(shape, dtype)):\n       with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch,\n                     threads=threads) as (bx, by, bz):\n           Q_shared = T.alloc_shared([block_M, dim], dtype)\n           K_shared = T.alloc_shared([block_N, dim], dtype)\n           V_shared = T.alloc_shared([block_N, dim], dtype)\n           acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)\n           acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)\n           acc_o = T.alloc_fragment([block_M, dim], accum_dtype)\n           m_prev = T.alloc_fragment([block_M], accum_dtype)\n           m_cur = T.alloc_fragment([block_M], accum_dtype)\n           alpha = T.alloc_fragment([block_M], accum_dtype)\n           p_sum = T.alloc_fragment([block_M], accum_dtype)\n           logsum = T.alloc_fragment([block_M], accum_dtype)\n           T.copy(Q[bz, bx * block_M:(bx + 1) * block_M, by, :], Q_shared)\n           T.fill(acc_o, 0)\n           T.fill(logsum, 0)\n           T.fill(m_cur, NEG)\n           loop_range = (T.ceildiv((bx + 1) * block_M, block_N)\n                         if is_causal else T.ceildiv(seq_len, block_N))\n           for k in T.Pipelined(loop_range, num_stages=num_stages):\n               T.copy(K[bz, k * block_N:(k + 1) * block_N, by, :], K_shared)\n               if is_causal:\n                   for i, j in T.Parallel(block_M, block_N):\n                       acc_s[i, j] = T.if_then_else(\n                           bx * block_M + i >= k * block_N + j, 0.0, NEG)\n               else:\n                   T.clear(acc_s)\n               T.gemm(Q_shared, K_shared, acc_s, transpose_B=True)\n               T.copy(V[bz, k * block_N:(k + 1) * block_N, by, :], V_shared)\n               T.copy(m_cur, m_prev)\n               T.reduce_max(acc_s, m_cur, dim=1, clear=False)\n               for i in T.Parallel(block_M):\n                   alpha[i] = T.exp((m_prev[i] - m_cur[i]) * scale)\n               for i, j in T.Parallel(block_M, block_N):\n                   acc_s[i, j] = T.exp((acc_s[i, j] - m_cur[i]) * scale)\n               T.reduce_sum(acc_s, p_sum, dim=1)\n               for i in T.Parallel(block_M):\n                   logsum[i] = logsum[i] * alpha[i] + p_sum[i]\n               for i, j in T.Parallel(block_M, dim):\n                   acc_o[i, j] = acc_o[i, j] * alpha[i]\n               T.copy(acc_s, acc_s_cast)\n               T.gemm(acc_s_cast, V_shared, acc_o)\n           for i, j in T.Parallel(block_M, dim):\n               acc_o[i, j] = acc_o[i, j] / logsum[i]\n           T.copy(acc_o, O[bz, bx * block_M:(bx + 1) * block_M, by, :])\n   return main\ndef section_6():\n   banner(\"6. FLASHATTENTION — fused attention forward\")\n   B, H, S, D = 4, 8, 1024, 64\n   q = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16)\n   k = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16)\n   v = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16)\n   def torch_ref(causal: bool):\n       qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v))\n       o = F.scaled_dot_product_attention(qt, kt, vt, is_causal=causal)\n       return o.transpose(1, 2).contiguous()\n   for causal in (False, True):\n       tag = \"causal\" if causal else \"full\"\n       attn = make_flash_attn(B, H, S, D, is_causal=causal,\n                              block_M=64, block_N=64,\n                              num_stages=1 if SM < 80 else 2, threads=128)\n       o = attn(q, k, v)\n       ref = torch_ref(causal)\n       check(o, ref, f\"flash_attn ({tag})\", tol=3e-2)\n       flops = 4 * B * H * S * S * D * (0.5 if causal else 1.0)\n       ms = bench(lambda: attn(q, k, v), warmup=5, rep=20)\n       ms_ref = bench(lambda: torch_ref(causal), warmup=5, rep=20)\n       print(f\"   {tag:>6}: tilelang {ms:6.3f} ms \"\n             f\"({flops/(ms*1e-3)/1e12:5.2f} TFLOP/s)   \"\n             f\"torch SDPA {ms_ref:6.3f} ms \"\n             f\"({flops/(ms_ref*1e-3)/1e12:5.2f} TFLOP/s)\")\n   print(\"   ~70 lines of Python for a fused, causal, tensor-core attention.\")\n   print(\"   The upstream repo pushes the same structure to FlashMLA-level perf\")\n   print(\"   on H100 with warp specialisation and TMA.\")\n```\n\nWe implement a fused FlashAttention forward kernel that processes query, key, and value tiles without materializing the full attention-score matrix in global memory. We apply online softmax updates using running maxima, normalization sums, rescaling factors, and tiled tensor-core matrix multiplications. We validate both causal and non-causal attention against PyTorch scaled dot-product attention and compare their latency and computational throughput.\n\n``` python\ndef section_7():\n   banner(\"7. AUTOTUNING — @tilelang.autotune\")\n   M = N = K = 2048\n   def configs():\n       out = []\n       for bm in (64, 128):\n           for bn in (128, 256):\n               for bk in (32, 64):\n                   for stages in (2, DEFAULT_STAGES):\n                       for thr in (128, 256):\n                           if smem_bytes(bm, bn, bk, stages) > SMEM_CAP:\n                               continue\n                           cfg = dict(block_M=bm, block_N=bn, block_K=bk,\n                                      num_stages=stages, threads=thr)\n                           if cfg not in out:\n                               out.append(cfg)\n       return out[:8]\n   space = configs()\n   print(f\"   searching {len(space)} configurations \"\n         f\"(each one is a real nvcc compile + benchmark)\")\n   @tilelang.autotune(configs=space, warmup=10, rep=20)\n   @tilelang.jit(out_idx=[-1])\n   def tuned_matmul(M: int, N: int, K: int,\n                    block_M: int = 128, block_N: int = 128, block_K: int = 32,\n                    num_stages: int = 2, threads: int = 128,\n                    dtype: str = \"float16\", accum_dtype: str = \"float\"):\n       @T.prim_func\n       def main(A: T.Tensor((M, K), dtype),\n                B: T.Tensor((K, N), dtype),\n                C: T.Tensor((M, N), dtype)):\n           with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M),\n                         threads=threads) as (bx, by):\n               A_shared = T.alloc_shared((block_M, block_K), dtype)\n               B_shared = T.alloc_shared((block_K, block_N), dtype)\n               C_local = T.alloc_fragment((block_M, block_N), accum_dtype)\n               T.clear(C_local)\n               for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):\n                   T.copy(A[by * block_M, ko * block_K], A_shared)\n                   T.copy(B[ko * block_K, bx * block_N], B_shared)\n                   T.gemm(A_shared, B_shared, C_local)\n               T.copy(C_local, C[by * block_M, bx * block_N])\n       return main\n   a = torch.randn(M, K, device=DEV, dtype=torch.float16)\n   b = torch.randn(K, N, device=DEV, dtype=torch.float16)\n   t0 = time.time()\n   from tilelang.autotuner import set_autotune_inputs\n   with set_autotune_inputs(a, b):\n       best = tuned_matmul(M, N, K)\n   print(f\"   tuning took {time.time()-t0:.0f} s\")\n   c = best(a, b)\n   check(c, a @ b, \"autotuned matmul\")\n   ms = bench(lambda: best(a, b))\n   print(f\"   best kernel: {ms:.3f} ms -> {2*M*N*K/(ms*1e-3)/1e12:.2f} TFLOP/s\")\n   print(\"   Re-running this cell hits the on-disk autotuner cache and is instant.\")\n   print(\"   Turn caching off with TILELANG_AUTO_TUNING_DISABLE_CACHE=1.\")\n```\n\nWe define an autotuning search space across matrix tile sizes, K-block dimensions, pipeline depths, and thread counts while filtering configurations that exceed the shared-memory budget. We use TileLang’s autotuning decorator to compile, benchmark, validate, and cache multiple kernel schedules for the same matrix-multiplication workload. We then execute the selected kernel, verify its output against PyTorch, and report the achieved latency and tensor-core throughput.\n\n``` python\n@tilelang.jit\ndef make_print_demo(N: int = 128, dtype: str = \"float32\"):\n   \"\"\"T.print emits a guarded device-side printf. Keep the grid tiny.\"\"\"\n   @T.prim_func\n   def main(A: T.Tensor((N,), dtype), B: T.Tensor((N,), dtype)):\n       with T.Kernel(1, threads=32) as bx:\n           A_local = T.alloc_fragment((N,), dtype)\n           T.copy(A, A_local)\n           for i in T.Parallel(N):\n               A_local[i] = A_local[i] * 2.0\n           T.print(A_local[0], msg=\"A_local[0] after doubling\")\n           T.copy(A_local, B)\n   return main\ndef section_8():\n   banner(\"8. INTROSPECTION — T.print, sources, profiler\")\n   try:\n       dbg = make_print_demo(128)\n       a = torch.ones(128, device=DEV)\n       b = torch.empty(128, device=DEV)\n       dbg(a, b)\n       torch.cuda.synchronize()\n       print(f\"   T.print fired above; host check: b[0] = {b[0].item()} (expect 2.0)\")\n   except Exception as e:\n       print(f\"   T.print demo skipped: {type(e).__name__}: {e}\")\n   kern = make_matmul(1024, 1024, 1024, 128, 128, 32,\n                      num_stages=min(2, DEFAULT_STAGES), threads=128)\n   src = kern.get_kernel_source()\n   print(f\"\\n   get_kernel_source(): {len(src.splitlines())} lines of CUDA\")\n   print(\"   grep-worthy landmarks:\")\n   for needle in (\"__global__\", \"extern \\\"C\\\"\", \"mma\", \"cp.async\",\n                  \"__syncthreads\", \"ldmatrix\"):\n       n = src.count(needle)\n       if n:\n           print(f\"      {needle:<16} x{n}\")\n   try:\n       prof = kern.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)\n       print(f\"\\n   profiler.do_bench(): {prof.do_bench():.3f} ms \"\n             f\"(it synthesises its own inputs)\")\n   except Exception as e:\n       print(f\"\\n   profiler skipped: {type(e).__name__}: {e}\")\n   print(\"\\n   Other tools worth knowing:\")\n   print(\"      kernel.get_host_source()     - the launch wrapper\")\n   print(\"      T.device_assert(cond, msg)   - device-side assertions\")\n   print(\"      TILELANG_DISABLE_CACHE=1     - force recompilation\")\n   print(\"      tilelang.tools.plot_layout   - visualise fragment/shared layouts\")\n   print(\"      env TILELANG_CACHE_DIR       - where compiled kernels live\")\nCHEATSHEET = \"\"\"\n  SCOPES              T.alloc_shared(shape, dtype)      __shared__ tile\n                      T.alloc_fragment(shape, dtype)    registers, layout inferred\n                      T.alloc_var(dtype)                one scalar per thread\n                      T.alloc_barrier(n)                mbarrier (Hopper pipelines)\n  MOVEMENT            T.copy(src, dst)                  vectorized + casting move\n                      T.async_copy(src, dst)            explicit cp.async\n                      T.tma_copy(...)                   Hopper bulk async\n                      T.transpose(src, dst)             shared-memory transpose\n  COMPUTE             T.gemm(A_s, B_s, C_f,\n                             transpose_A/B=...,\n                             policy=T.GemmWarpPolicy.*) tile matmul on tensor cores\n                      T.gemm_sp(...)                    2:4 structured sparsity\n                      T.reduce_sum/max/min(buf, out, dim=, clear=)\n                      T.cumsum / T.cummax               scans\n                      T.clear(buf) / T.fill(buf, v)\n  LOOPS               T.Parallel(*extents)              elementwise -> threads\n                      T.Pipelined(n, num_stages=k)      software pipeline\n                      T.serial(n)                       plain sequential loop\n  ANNOTATIONS         T.use_swizzle(panel_size=, enable=)   L2 rasterization\n                      T.annotate_layout({buf: layout})      explicit layouts\n                      T.annotate_l2_hit_ratio(buf, r)\n  ENTRY POINTS        @tilelang.jit(out_idx=[-1])       last arg is the return value\n                      @tilelang.autotune(configs=...)   stack above @jit\n                      kernel.get_kernel_source()\n                      kernel.get_profiler().do_bench()\n  ENV VARS            TILELANG_CACHE_DIR, TILELANG_DISABLE_CACHE,\n                      TILELANG_AUTO_TUNING_DISABLE_CACHE,\n                      TILELANG_AUTO_TUNING_MAX_CPU_COUNT\n  WHERE TO GO NEXT    examples/flash_attention        fwd + bwd, autotuned\n                      examples/dequantize_gemm        W4A16 with LOP3 tricks\n                      examples/deepseek_mla           MLA decode, ~80 lines\n                      examples/linear_attention       RetNet / Mamba\n                      github.com/tile-ai/tilelang-puzzles   10 graded exercises\n                      tilelang.com                    full docs + API reference\n\"\"\"\nSECTIONS = [\n   (\"1  hello / vector add\", section_1),\n   (\"2  tiled GEMM\", section_2),\n   (\"3  schedule sweep\", section_3),\n   (\"4  epilogue fusion\", section_4),\n   (\"5  softmax reduction\", section_5),\n   (\"6  flash attention\", section_6),\n   (\"7  autotuning\", section_7),\n   (\"8  introspection\", section_8),\n]\ndef main(only=None):\n   \"\"\"only: optional list of 1-based section numbers, e.g. main([2, 6]).\"\"\"\n   t_start = time.time()\n   status = []\n   for idx, (name, fn) in enumerate(SECTIONS, start=1):\n       if only and idx not in only:\n           continue\n       t0 = time.time()\n       try:\n           fn()\n           status.append((name, \"ok\", time.time() - t0))\n       except Exception:\n           print(\"\\n   !! section failed — continuing with the rest\\n\")\n           traceback.print_exc()\n           status.append((name, \"FAILED\", time.time() - t0))\n       finally:\n           torch.cuda.synchronize()\n   banner(\"9. SUMMARY & CHEATSHEET\")\n   for name, st, dt in status:\n       print(f\"   {st:>6}  {name:<24} {dt:6.1f} s\")\n   print(f\"\\n   total: {time.time()-t_start:.0f} s\")\n   print(CHEATSHEET)\nif __name__ == \"__main__\":\n   main()\n```\n\nWe introduce TileLang’s debugging and introspection workflow through device-side printing, generated CUDA inspection, and the built-in kernel profiler. We examine compiler-emitted landmarks such as tensor-core operations, asynchronous copies, synchronization barriers, and matrix-load instructions. We finally organize all tutorial sections into a fault-tolerant runner that records execution status, reports timing information, and prints a compact TileLang programming reference.\n\nIn conclusion, we built a practical understanding of how TileLang translates tile-level Python programs into optimized GPU kernels without requiring us to manually manage thread indices, warp-level data layouts, tensor-core instructions, or asynchronous memory barriers. We implemented and validated kernels that cover bandwidth-bound elementwise operations, compute-intensive GEMM workloads, fused neural-network epilogues, register-resident reductions, and online-softmax attention. We also examined how block dimensions, shared-memory consumption, pipeline depth, thread count, tile shape, and L2 swizzling influence performance across different GPU architectures. Finally, we used generated-source inspection, device-side debugging, profiling, and automated schedule search to establish a complete workflow for developing, verifying, benchmarking, and refining custom TileLang kernels.\n\nCheck out the ** Full Code here. **Also, feel free to follow us on\n\n**and don’t forget to join our**[Twitter](https://x.com/intent/follow?screen_name=marktechpost)\n\n**and Subscribe to**\n\n[150k+ML SubReddit](https://www.reddit.com/r/machinelearningnews/)**. Wait! are you on telegram?**\n\n[our Newsletter](https://www.aidevsignals.com/)\n\n[now you can join us on telegram as well.](https://t.me/machinelearningresearchnews)Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? [Connect with us](https://forms.gle/wbash1wF6efRj8G58)\n\nSana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.", "url": "https://wpnews.pro/news/designing-high-performance-gpu-kernels-with-tilelang-tensor-core-gemm-fused-and", "canonical_source": "https://www.marktechpost.com/2026/07/25/designing-high-performance-gpu-kernels-with-tilelang-tensor-core-gemm-fused-softmax-flashattention-and-autotuning/", "published_at": "2026-07-25 18:08:12+00:00", "updated_at": "2026-07-25 18:37:30.006776+00:00", "lang": "en", "topics": ["developer-tools", "machine-learning", "artificial-intelligence"], "entities": ["TileLang", "TVM", "PyTorch", "cuBLAS", "CUDA"], "alternates": {"html": "https://wpnews.pro/news/designing-high-performance-gpu-kernels-with-tilelang-tensor-core-gemm-fused-and", "markdown": "https://wpnews.pro/news/designing-high-performance-gpu-kernels-with-tilelang-tensor-core-gemm-fused-and.md", "text": "https://wpnews.pro/news/designing-high-performance-gpu-kernels-with-tilelang-tensor-core-gemm-fused-and.txt", "jsonld": "https://wpnews.pro/news/designing-high-performance-gpu-kernels-with-tilelang-tensor-core-gemm-fused-and.jsonld"}}