Skip to content

[Feature] Introduce tile scheduler #2441

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LeiWang1999 merged 2 commits into from
Jun 24, 2026
Merged

[Feature] Introduce tile scheduler #2441

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
958430b
[Feature] Implement PersistentTileScheduler for warp-specialized GEMM
Rachmanino Jun 23, 2026
6874615
[Refactor] Update PersistentTileScheduler usage in GEMM examples
Rachmanino Jun 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
307 changes: 153 additions & 154 deletions examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,11 @@
# Persistent, num_epi_stages = 2
# Persistent, warp-specialized SM100 GEMM using PersistentTileScheduler.
#
# Each warp role (TMA / tcgen5 MMA / epilogue) creates its own scheduler instance
# and drives a single ``while sched.valid()`` loop. The scheduler owns the only
# iteration clock: ``sched.current_iter[0]`` is the wave index ``w`` (used for the
# pipeline phase / ``w & 1`` double-buffering) and ``sched.m_idx[0]`` /
# ``sched.n_idx[0]`` are the tile coords. One clock, held by the scheduler -- no
# separate ``for w in range(waves)`` counter.

import torch
import tilelang
Expand Down Expand Up @@ -32,7 +39,6 @@ def gemm_persistent(
n_blocks = T.ceildiv(N, block_N)
assert K % (2 * block_K) == 0 # for simplicity
k_blocks = T.ceildiv(K, block_K)
waves = T.ceildiv(m_blocks * n_blocks, sm_num)
group_size = 8
assert n_blocks % (2 * group_size) == 0 # Please adjust group_size if not satisfied

Expand All @@ -52,81 +58,77 @@ def gemm_persistent(
tx = T.get_thread_binding()

if tx < 32: # warp 0: issue tma
for w in T.unroll(waves):
tile_id = sm_num * w + block_id
bx = (tile_id // group_size) % m_blocks
by = (tile_id % group_size) + (tile_id // group_size) // m_blocks * group_size

if bx * block_M < M and by * block_N < N:
for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(consumed[phase % num_stages], ((phase // num_stages) & 1) ^ 1)
T.tma_copy(
A[bx * block_M : (bx + 1) * block_M, k * block_K : (k + 1) * block_K],
A_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
)
T.tma_copy(
B[k * block_K : (k + 1) * block_K, by * block_N : (by + 1) * block_N],
B_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
)
T.mbarrier_arrive(loaded[phase % num_stages])
sched = T.PersistentTileScheduler("sched_tma", m_blocks, n_blocks, swizzle_size=group_size)
sched.init(block_id)
while sched.valid():
w, bx, by = sched.current_iter[0], sched.m_idx[0], sched.n_idx[0]

for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(consumed[phase % num_stages], ((phase // num_stages) & 1) ^ 1)
T.tma_copy(
A[bx * block_M : (bx + 1) * block_M, k * block_K : (k + 1) * block_K],
A_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
)
T.tma_copy(
B[k * block_K : (k + 1) * block_K, by * block_N : (by + 1) * block_N],
B_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
)
T.mbarrier_arrive(loaded[phase % num_stages])
sched.next_tile()

elif tx < 64: # warp 1: issue tcgen5
for w in T.unroll(waves):
tile_id = sm_num * w + block_id
bx = (tile_id // group_size) % m_blocks
by = (tile_id % group_size) + (tile_id // group_size) // m_blocks * group_size

if bx * block_M < M and by * block_N < N:
T.mbarrier_wait_parity(tmem_empty[w & 1], ((w // 2) & 1) ^ 1)
for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(loaded[phase % num_stages], (phase // num_stages) & 1)
if w & 1 == 0:
T.tcgen05_gemm(
A_shared[k % num_stages, :, :],
B_shared[k % num_stages, :, :],
C_tmem_0,
False,
False,
mbar=consumed[k % num_stages],
clear_accum=k == 0,
)
else:
T.tcgen05_gemm(
A_shared[k % num_stages, :, :],
B_shared[k % num_stages, :, :],
C_tmem_1,
False,
False,
mbar=consumed[k % num_stages],
clear_accum=k == 0,
)
T.tcgen05_mma_arrive(tmem_full[w & 1])

elif 128 <= tx < 256: # warp 4~7: epilogue
for w in T.unroll(waves):
tile_id = sm_num * w + block_id
bx = (tile_id // group_size) % m_blocks
by = (tile_id % group_size) + (tile_id // group_size) // m_blocks * group_size

if bx * block_M < M and by * block_N < N:
T.mbarrier_wait_parity(tmem_full[w & 1], (w // 2) & 1)
if (w & 1) == 0:
T.copy(C_tmem_0, C_local)
sched = T.PersistentTileScheduler("sched_mma", m_blocks, n_blocks, swizzle_size=group_size)
sched.init(block_id)
while sched.valid():
w = sched.current_iter[0]

T.mbarrier_wait_parity(tmem_empty[w & 1], ((w // 2) & 1) ^ 1)
for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(loaded[phase % num_stages], (phase // num_stages) & 1)
if w & 1 == 0:
T.tcgen05_gemm(
A_shared[k % num_stages, :, :],
B_shared[k % num_stages, :, :],
C_tmem_0,
mbar=consumed[k % num_stages],
clear_accum=k == 0,
)
else:
T.copy(C_tmem_1, C_local)
T.mbarrier_arrive(tmem_empty[w & 1])
T.tcgen05_gemm(
A_shared[k % num_stages, :, :],
B_shared[k % num_stages, :, :],
C_tmem_1,
mbar=consumed[k % num_stages],
clear_accum=k == 0,
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved
T.tcgen05_mma_arrive(tmem_full[w & 1])
sched.next_tile()

if use_tma_store:
for i in T.unroll(T.ceildiv(block_N, store_block_N)):
T.copy(C_local[:, i * store_block_N : (i + 1) * store_block_N], C_shared)
T.copy(C_shared, C[bx * block_M, by * block_N + i * store_block_N])
else:
T.copy(C_local, C_local_cast)
T.copy(C_local_cast, C[bx * block_M, by * block_N])
elif 128 <= tx < 256: # warp 4~7: epilogue
sched = T.PersistentTileScheduler("sched_epi", m_blocks, n_blocks, swizzle_size=group_size)
sched.init(block_id)
while sched.valid():
w, bx, by = sched.current_iter[0], sched.m_idx[0], sched.n_idx[0]

T.mbarrier_wait_parity(tmem_full[w & 1], (w // 2) & 1)
if (w & 1) == 0:
T.copy(C_tmem_0, C_local)
else:
T.copy(C_tmem_1, C_local)
T.mbarrier_arrive(tmem_empty[w & 1])

if use_tma_store:
for i in T.unroll(T.ceildiv(block_N, store_block_N)):
T.copy(C_local[:, i * store_block_N : (i + 1) * store_block_N], C_shared)
T.copy(C_shared, C[bx * block_M, by * block_N + i * store_block_N])
else:
T.copy(C_local, C_local_cast)
T.copy(C_local_cast, C[bx * block_M, by * block_N])
sched.next_tile()
return C


Expand All @@ -151,15 +153,13 @@ def gemm_persistent_2cta(
C = T.empty((M, N), out_dtype)

sm_num = driver.get_num_sms()
num_clusters = sm_num // 2
m_blocks = T.ceildiv(M, block_M)
m_clusters = m_blocks // 2
n_blocks = T.ceildiv(N, block_N)
assert K % (2 * block_K) == 0 # for simplicity
k_blocks = T.ceildiv(K, block_K)
waves = T.ceildiv(m_blocks * n_blocks, sm_num)
group_size = 8 # in cluster
assert n_blocks % (2 * group_size) == 0 # Please adjust group_size if not satisfied
cluster_size = 2

with T.ClusterKernel(sm_num, threads=256, cluster_dims=2) as (block_id):
A_shared = T.alloc_shared((num_stages, block_M, block_K), in_dtype)
Expand All @@ -179,89 +179,85 @@ def gemm_persistent_2cta(
T.assume(cta_id < 2) # todo: automatically assume this

if tx < 32: # warp 0: issue tma
for w in T.unroll(waves):
# manual threadblock swizzle
cluster_id = block_id // 2
tile_id = num_clusters * w + cluster_id
bx_cluster = (tile_id // group_size) % m_clusters
bx = bx_cluster * 2 + cta_id
by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size

if bx * block_M < M and by * block_N < N:
for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(consumed[phase % num_stages], ((phase // num_stages) & 1) ^ 1)
T.tma_copy(
A[bx * block_M : (bx + 1) * block_M, k * block_K : (k + 1) * block_K],
sched = T.PersistentTileScheduler(
"sched_tma", m_blocks, n_blocks, swizzle_size=group_size, cluster_size=cluster_size)
sched.init(block_id // cluster_size)
while sched.valid():
w = sched.current_iter[0]
bx, by = sched.m_idx[0] * cluster_size + cta_id, sched.n_idx[0]

for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(consumed[phase % num_stages], ((phase // num_stages) & 1) ^ 1)
T.tma_copy(
A[bx * block_M : (bx + 1) * block_M, k * block_K : (k + 1) * block_K],
A_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
)

T.tma_copy(
B[k * block_K : (k + 1) * block_K, (by * 2 + cta_id) * block_N // 2 : (by * 2 + cta_id + 1) * block_N // 2],
B_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
)
T.mbarrier_arrive(loaded[phase % num_stages], 0)
sched.next_tile()

elif tx < 64 and cta_id == 0: # warp 1: issue tcgen5
sched = T.PersistentTileScheduler(
"sched_mma", m_blocks, n_blocks, swizzle_size=group_size, cluster_size=cluster_size)
sched.init(block_id // cluster_size)
while sched.valid():
w = sched.current_iter[0]

T.mbarrier_wait_parity(tmem_empty[w & 1], ((w // 2) & 1) ^ 1)
for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(loaded[phase % num_stages], (phase // num_stages) & 1)
if w & 1 == 0:
T.tcgen05_gemm(
A_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
B_shared[phase % num_stages, :, :],
C_tmem_0,
mbar=consumed[phase % num_stages],
clear_accum=k == 0,
use_2cta=True,
)

T.tma_copy(
B[k * block_K : (k + 1) * block_K, (by * 2 + cta_id) * block_N // 2 : (by * 2 + cta_id + 1) * block_N // 2],
else:
T.tcgen05_gemm(
A_shared[phase % num_stages, :, :],
B_shared[phase % num_stages, :, :],
barrier=loaded[phase % num_stages],
C_tmem_1,
mbar=consumed[phase % num_stages],
clear_accum=k == 0,
use_2cta=True,
)
T.mbarrier_arrive(loaded[phase % num_stages], 0)

elif tx < 64 and cta_id == 0: # warp 1: issue tcgen5
for w in T.unroll(waves):
# manual threadblock swizzle
cluster_id = block_id // 2
tile_id = num_clusters * w + cluster_id
bx_cluster = (tile_id // group_size) % m_clusters
bx = bx_cluster * 2 + cta_id
by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size

if bx * block_M < M and by * block_N < N:
T.mbarrier_wait_parity(tmem_empty[w & 1], ((w // 2) & 1) ^ 1)
for k in T.serial(k_blocks):
phase = w * k_blocks + k
T.mbarrier_wait_parity(loaded[phase % num_stages], (phase // num_stages) & 1)
if w & 1 == 0:
T.tcgen05_gemm(
A_shared[phase % num_stages, :, :],
B_shared[phase % num_stages, :, :],
C_tmem_0,
mbar=consumed[phase % num_stages],
clear_accum=k == 0,
use_2cta=True,
)
else:
T.tcgen05_gemm(
A_shared[phase % num_stages, :, :],
B_shared[phase % num_stages, :, :],
C_tmem_1,
mbar=consumed[phase % num_stages],
clear_accum=k == 0,
use_2cta=True,
)
T.tcgen05_mma_arrive(tmem_full[w & 1], arrive_2cta=True)
T.tcgen05_mma_arrive(tmem_full[w & 1], arrive_2cta=True)
sched.next_tile()

elif 128 <= tx < 256: # warp 4~7: epilogue
for w in T.unroll(waves):
# manual threadblock swizzle
cluster_id = block_id // 2
tile_id = num_clusters * w + cluster_id
bx_cluster = (tile_id // group_size) % m_clusters
bx = bx_cluster * 2 + cta_id
by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size

if bx * block_M < M and by * block_N < N:
T.mbarrier_wait_parity(tmem_full[w & 1], (w // 2) & 1)
if (w & 1) == 0:
T.copy(C_tmem_0, C_local)
else:
T.copy(C_tmem_1, C_local)
T.mbarrier_arrive(tmem_empty[w & 1], 0)

if use_tma_store:
for i in T.unroll(T.ceildiv(block_N, store_block_N)):
T.copy(C_local[:, i * store_block_N : (i + 1) * store_block_N], C_shared)
T.copy(C_shared, C[bx * block_M, by * block_N + i * store_block_N])
else:
T.copy(C_local, C_local_cast)
T.copy(C_local_cast, C[bx * block_M, by * block_N])
sched = T.PersistentTileScheduler(
"sched_epi", m_blocks, n_blocks, swizzle_size=group_size, cluster_size=cluster_size)
sched.init(block_id // cluster_size)
while sched.valid():
w = sched.current_iter[0]
bx, by = sched.m_idx[0] * cluster_size + cta_id, sched.n_idx[0]

T.mbarrier_wait_parity(tmem_full[w & 1], (w // 2) & 1)
if (w & 1) == 0:
T.copy(C_tmem_0, C_local)
else:
T.copy(C_tmem_1, C_local)
T.mbarrier_arrive(tmem_empty[w & 1], 0)

if use_tma_store:
for i in T.unroll(T.ceildiv(block_N, store_block_N)):
T.copy(C_local[:, i * store_block_N : (i + 1) * store_block_N], C_shared)
T.copy(C_shared, C[bx * block_M, by * block_N + i * store_block_N])
else:
T.copy(C_local, C_local_cast)
T.copy(C_local_cast, C[bx * block_M, by * block_N])
sched.next_tile()

return C

Expand All @@ -285,7 +281,10 @@ def main():
print("All checks passed. ✅")

tl_latency = do_bench(
lambda: kernel(a, b, block_M, block_N, store_block_N, block_K, in_dtype, out_dtype, accum_dtype, num_stages), backend="cupti"
lambda: kernel(a, b, block_M, block_N, store_block_N, block_K, in_dtype, out_dtype, accum_dtype, num_stages),
_n_warmup=50,
_n_repeat=50,
backend="cupti"
)
torch_latency = do_bench(lambda: a @ b, backend="cupti")
print(f"Tilelang latency: {tl_latency} ms")
Expand Down
10 changes: 10 additions & 0 deletions tilelang/language/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,16 @@
clc_get_first_ctaid_z, # noqa: F401
)

from .meta import (
inline, # noqa: F401
meta_class, # noqa: F401
)

from .tile_schedule import (
BaseTileScheduler, # noqa: F401
PersistentTileScheduler, # noqa: F401
)


def import_source(source: str | None = None):
# source is the source code to be imported
Expand Down
Loading
Loading