TileLang Patterns

Canonical TileLang code patterns from official examples. Source: tile-ai/tilelang/examples

1. Element-wise: Load → Compute → Store

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_N), dtype)
    B_shared = T.alloc_shared((BLOCK_M, BLOCK_N), dtype)
    C_local  = T.alloc_fragment((BLOCK_M, BLOCK_N), out_dtype)
    C_shared = T.alloc_shared((BLOCK_M, BLOCK_N), out_dtype)

    T.copy(A[by * BLOCK_M, bx * BLOCK_N], A_shared)
    T.copy(B[by * BLOCK_M, bx * BLOCK_N], B_shared)
    for i, j in T.Parallel(BLOCK_M, BLOCK_N):
        C_local[i, j] = A_shared[i, j] + B_shared[i, j]
    T.copy(C_local, C_shared)
    T.copy(C_shared, C[by * BLOCK_M, bx * BLOCK_N])

2. Reduction with Serial Accumulation

with T.Kernel(N // BLOCK_N, threads=256) as pid:
    A_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
    acc     = T.alloc_fragment((BLOCK_N,), accum_dtype)
    T.clear(acc)
    for m in T.Serial(M // BLOCK_M):
        T.copy(A[pid * BLOCK_N, m * BLOCK_M], A_local)
        T.reduce_sum(A_local, acc, dim=1, clear=False)   # accumulate
    T.copy(acc, B[pid * BLOCK_N])

3. GEMM: Shared Memory + Pipelined Loop

with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) 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), T.float32)
    T.clear(C_local)
    for k in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
        T.copy(A[by * BLOCK_M, k * BLOCK_K], A_shared)
        T.copy(B[k * 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])

Typical configs:

SM80: block_M=128, block_N=256, block_K=32, num_stages=2, threads=128
SM90: block_M=128, block_N=256, block_K=64, num_stages=3, threads=256

4. GEMM with Rasterization (L2 Swizzle)

Same as pattern 3, with T.use_swizzle for L2 cache locality:

with T.Kernel(...) as (bx, by):
    T.use_swizzle(panel_size=10)
    # ... same GEMM body ...

5. GEMM with Fragment → Shared → Global Store

When output needs type conversion through shared memory:

    C_shared = T.alloc_shared((BLOCK_M, BLOCK_N), out_dtype)
    # ... GEMM loop ...
    T.copy(C_local, C_shared)                          # fragment → shared
    T.copy(C_shared, C[by * BLOCK_M, bx * BLOCK_N])   # shared → global

6. Persistent GEMM (Wave-level Scheduling)

with T.Kernel(sm_num, threads=256) as block_id:
    C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), T.float32)
    for bx, by in T.Persistent([T.ceildiv(M, BLOCK_M), T.ceildiv(N, BLOCK_N)],
                                 sm_num, block_id):
        T.clear(C_local)
        for k in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=num_stages):
            T.copy(A[by * BLOCK_M, k * BLOCK_K], A_shared)
            T.copy(B[k * 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])

7. GEMV: Element-wise Multiply + Reduce

with T.Kernel(T.ceildiv(M, BLOCK_M), threads=128) as pid_m:
    A_local  = T.alloc_fragment((BLOCK_M, BLOCK_K), dtype)
    B_local  = T.alloc_fragment((BLOCK_K,), dtype)
    C_local  = T.alloc_fragment((BLOCK_M,), T.float32)
    AB_temp  = T.alloc_fragment((BLOCK_M, BLOCK_K), T.float32)
    T.clear(C_local)
    for k in T.Serial(K // BLOCK_K):
        T.copy(A[pid_m * BLOCK_M, k * BLOCK_K], A_local)
        T.copy(B[k * BLOCK_K,], B_local)
        for i, j in T.Parallel(BLOCK_M, BLOCK_K):
            AB_temp[i, j] = A_local[i, j].astype(T.float32) * B_local[j].astype(T.float32)
        T.reduce_sum(AB_temp, C_local, dim=1, clear=False)
    T.copy(C_local, C[pid_m * BLOCK_M,])

8. Flash Attention Forward (Online Softmax)

with T.Kernel(T.ceildiv(seq_len, BLOCK_M), heads, batch, threads=128) as (bx, by, bz):
    Q_shared = T.alloc_shared((BLOCK_M, dim), dtype)
    K_shared = T.alloc_shared((BLOCK_N, dim), dtype)
    V_shared = T.alloc_shared((BLOCK_N, dim), dtype)
    acc_s    = T.alloc_fragment((BLOCK_M, BLOCK_N), T.float32)
    acc_o    = T.alloc_fragment((BLOCK_M, dim), T.float32)
    scores_max      = T.alloc_fragment((BLOCK_M,), T.float32)
    scores_max_prev = T.alloc_fragment((BLOCK_M,), T.float32)
    scores_scale    = T.alloc_fragment((BLOCK_M,), T.float32)
    scores_sum      = T.alloc_fragment((BLOCK_M,), T.float32)
    logsum          = T.alloc_fragment((BLOCK_M,), T.float32)

    T.clear(acc_o)
    T.clear(logsum)
    T.fill(scores_max, -T.infinity(T.float32))
    T.copy(Q[bz, by, bx * BLOCK_M, :], Q_shared)

    loop_range = T.ceildiv((bx + 1) * BLOCK_M, BLOCK_N) if is_causal else T.ceildiv(seq_len, BLOCK_N)

    for k in T.Pipelined(loop_range, num_stages=1):
        T.copy(K[bz, by, k * BLOCK_N, :], K_shared)
        T.copy(V[bz, by, k * BLOCK_N, :], V_shared)

        # S = Q @ K^T
        T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, clear_accum=True)

        # Online softmax
        T.copy(scores_max, scores_max_prev)
        T.reduce_max(acc_s, scores_max, dim=1, clear=False)
        for i in T.Parallel(BLOCK_M):
            scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
        for i in T.Parallel(BLOCK_M):
            scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
        for i, j in T.Parallel(BLOCK_M, dim):
            acc_o[i, j] *= scores_scale[i]
        for i, j in T.Parallel(BLOCK_M, BLOCK_N):
            acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
        T.reduce_sum(acc_s, scores_sum, dim=1, clear=True)
        for i in T.Parallel(BLOCK_M):
            logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]

        # O += S @ V
        T.gemm(acc_s, V_shared, acc_o)

    # Normalize
    for i, j in T.Parallel(BLOCK_M, dim):
        acc_o[i, j] /= logsum[i]
    T.copy(acc_o, Output[bz, by, bx * BLOCK_M, :])

Key: scale = 1.44269504 (log2(e)) for T.exp2-based softmax.

9. Online Softmax (Standalone)

T.fill(lse, -T.infinity(accum_dtype))
for n in T.Pipelined(NN):
    T.copy(X[pid, n * BN], x_local)
    T.reduce_max(x_local, max_x, dim=0, clear=True)
    for j in T.Parallel(BN):
        exp_x[j] = T.exp2(x_local[j] * scale - max_x[0] * scale)
    T.reduce_sum(exp_x, sum_exp_x, dim=0, clear=True)
    lse[0] = max_x[0] * scale + T.log2(
        T.exp2(lse[0] - max_x[0] * scale) + sum_exp_x[0])

for n in T.Pipelined(NN):
    T.copy(X[pid, n * BN], x_local)
    for j in T.Parallel(BN):
        y_local[j] = T.exp2(x_local[j] * scale - lse[0])
    T.copy(y_local, Y[pid, n * BN])

10. INT4 Dequant Fused GEMM

B_shared  = T.alloc_shared((BLOCK_K, BLOCK_N // 2), T.uint8)
B_dequant = T.alloc_shared((BLOCK_K, BLOCK_N), T.float16)

for k in T.Pipelined(K // BLOCK_K, num_stages=num_stages):
    T.copy(A[...], A_shared)
    T.copy(B_packed[...], B_shared)
    for i, j in T.Parallel(BLOCK_K, BLOCK_N // 2):
        B_dequant[i, j * 2]     = T.cast(B_shared[i, j] & 0x0F, T.float16) - 8.0
        B_dequant[i, j * 2 + 1] = T.cast((B_shared[i, j] >> 4) & 0x0F, T.float16) - 8.0
    T.gemm(A_shared, B_dequant, C_local)

11. FP8 GEMM

with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
    A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), T.float8_e4m3fn)
    B_shared = T.alloc_shared((BLOCK_N, BLOCK_K), T.float8_e4m3fn)   # note: B is (N, K)
    C_local  = T.alloc_fragment((BLOCK_M, BLOCK_N), T.float32)
    T.clear(C_local)
    for k in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
        T.copy(A[by * BLOCK_M, k * BLOCK_K], A_shared)
        T.copy(B[bx * BLOCK_N, k * BLOCK_K], B_shared)
        T.gemm(A_shared, B_shared, C_local, transpose_B=True)   # B transposed
    T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N])

12. Convolution via im2col + GEMM

On Hopper (SM90):

T.c2d_im2col(img, col_shared, nhw_step, c_step, kernel, stride, dilation, pad)
T.gemm(col_shared, weight_shared, C_local)

On other GPUs (manual im2col):

for i, j in T.Parallel(BLOCK_M, BLOCK_K):
    m = by * BLOCK_M + i
    k = ko * BLOCK_K + j
    access_h = m % (OH * OW) // OW * S + k // (KW * C) * D - P
    access_w = m % OW * S + k // C % KW * D - P
    in_bound = (access_h >= 0) and (access_w >= 0) and (access_h < H) and (access_w < W)
    data_shared[i, j] = T.if_then_else(in_bound, data[n, access_h, access_w, k % C], 0)

13. Warp Specialization (Producer/Consumer)

with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=256) as (bx, by):
    data_is_ready   = T.alloc_barrier(arrive_count=128)
    compute_is_done = T.alloc_barrier(arrive_count=128)

    for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=0):
        with T.ws(0):   # warp group 0: producer (loads data)
            T.barrier_wait(compute_is_done, (ko + 1) % 2)
            T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared)
            T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared)
            T.barrier_arrive(data_is_ready)
        with T.ws(1):   # warp group 1: consumer (computes)
            T.barrier_wait(data_is_ready, ko % 2)
            T.gemm(A_shared, B_shared, C_local)
            T.barrier_arrive(compute_is_done)

    with T.ws(1):
        T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N])

14. RMS Normalization

with T.Kernel(T.ceildiv(M, BLOCK_M), threads=threads) as bx:
    A_shared    = T.alloc_shared((BLOCK_M, N), dtype)
    A_local     = T.alloc_fragment((BLOCK_M, N), T.float32)
    A_pow_local = T.alloc_fragment((BLOCK_M, N), T.float32)
    A_powsum    = T.alloc_fragment((BLOCK_M,), T.float32)

    T.copy(A[bx * BLOCK_M:, :], A_shared)
    T.copy(A_shared, A_local)
    for i, j in T.Parallel(BLOCK_M, N):
        A_pow_local[i, j] = A_local[i, j] * A_local[i, j]
    T.reduce_sum(A_pow_local, A_powsum, dim=1)
    for i in T.Parallel(BLOCK_M):
        A_powsum[i] = T.rsqrt(A_powsum[i] / N + 1e-12)
    for i, j in T.Parallel(BLOCK_M, N):
        A_local[i, j] *= A_powsum[i]
    T.copy(A_local, O_shared)
    T.copy(O_shared, O[bx * BLOCK_M:, :])

15. StreamK (Dynamic Work Distribution)

with T.Kernel(streamk_programs, threads=threads) as pid:
    # Compute start/end iterations for this SM
    start_iter = T.alloc_var(T.int32)
    end_iter   = T.alloc_var(T.int32)
    # ... compute bounds ...

    while start_iter[0] < last_iter:
        tile_id = start_iter[0] // iters_per_tile
        # ... compute partial tile ...
        if is_full_tile:
            T.copy(C_local, C[...])                    # direct write
        else:
            for i, j in T.Parallel(BLOCK_M, BLOCK_N):
                T.atomic_add(C[...], C_local[i, j])    # atomic for partials

16. Fused MoE (Gate + Up + Down)

# Kernel 1: Gate and Up projections fused
for k in T.Pipelined(T.ceildiv(dhidden, BLOCK_K), num_stages=num_stages):
    T.copy(input[...], input_shared)
    T.copy(W_gate[...], W_gate_shared)
    T.copy(W_up[...], W_up_shared)
    T.gemm(input_shared, W_gate_shared, gate_local, transpose_B=True)
    T.gemm(input_shared, W_up_shared, up_local, transpose_B=True)

# Fused SiLU + element-wise product
for i, j in T.Parallel(BLOCK_M, BLOCK_N):
    gate_local[i, j] = gate_local[i, j] * (1.0 / (1.0 + T.exp2(-gate_local[i, j] * 1.44269504)))
    up_local[i, j] *= gate_local[i, j]

17. 2D Grid with Broadcasting

with T.Kernel(N // BLOCK_N, M // BLOCK_M, threads=256) as (pid_n, pid_m):
    A_local = T.alloc_fragment((BLOCK_N,), dtype)
    B_local = T.alloc_fragment((BLOCK_M,), dtype)
    C_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
    T.copy(A[pid_n * BLOCK_N], A_local)
    T.copy(B[pid_m * BLOCK_M], B_local)
    for i, j in T.Parallel(BLOCK_N, BLOCK_M):
        C_local[i, j] = A_local[i] + B_local[j]
    T.copy(C_local, C[pid_n * BLOCK_N, pid_m * BLOCK_M])

18. Sparse GEMM (Structured 2:4 Sparsity)

A_sparse: T.Tensor((M, K // 2), dtype)    # pruned to 2:4 sparsity
E: T.Tensor(metadata_shape, T.uint8)       # sparsity metadata
B: T.Tensor((K, N), dtype)                 # dense

T.gemm_sp(A_shared, B_shared, C_local, E_shared)

Configuration Quick Reference

Kernel Type	threads	BLOCK_M	BLOCK_N	BLOCK_K	num_stages	Pass Configs
GEMM (SM80)	128	128	256	32	2	default
GEMM (SM90)	256	128	256	64	3	default
Flash Attn	128	64	64	—	1	`FAST_MATH`
Convolution	256	64	128	32	3	default
Linear Attn	128	varies	varies	varies	varies	`DISABLE_TMA`, `DISABLE_WARP_SPEC`
MoE	128-256	varies	varies	varies	2-3	`DISABLE_TMA`, `DISABLE_WARP_SPEC`
Elementwise	256	64	64	—	—	`DISABLE_TMA`, `DISABLE_WARP_SPEC`
RMS Norm	128-256	varies	N	—	—	`DISABLE_TMA`, `DISABLE_WARP_SPEC`

Weekly Reports

Weekly Report

Maintain Dashboard

Summary
- CI Health Log
- Ops Health Dashboard

AI Knowledge Base

TileLang Knowledge Base

Uh oh!

TileLang Patterns

TileLang Patterns

1. Element-wise: Load → Compute → Store

2. Reduction with Serial Accumulation

3. GEMM: Shared Memory + Pipelined Loop

4. GEMM with Rasterization (L2 Swizzle)

5. GEMM with Fragment → Shared → Global Store

6. Persistent GEMM (Wave-level Scheduling)

7. GEMV: Element-wise Multiply + Reduce

8. Flash Attention Forward (Online Softmax)

9. Online Softmax (Standalone)

10. INT4 Dequant Fused GEMM

11. FP8 GEMM

12. Convolution via im2col + GEMM

13. Warp Specialization (Producer/Consumer)

14. RMS Normalization

15. StreamK (Dynamic Work Distribution)

16. Fused MoE (Gate + Up + Down)

17. 2D Grid with Broadcasting

18. Sparse GEMM (Structured 2:4 Sparsity)

Configuration Quick Reference

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Clone this wiki locally