[Feature] 2-SM support for TMA, TMEM and TCGEN5MMA on Blackwell - #1882

Merged

LeiWang1999 merged 39 commits into

tile-ai:mainfrom

Rachmanino:wt/2sm

Mar 24, 2026

Commits on Mar 18, 2026

support 2-cta alloc, dealloc, umma_arrive
Rachmanino
committed
9f9677a
draft tma load 2sm support
Rachmanino
committed
fbc6bec
draft tcgen5 2sm support
Rachmanino
committed
4dc2901
add draft lower_blackwell_2sm pass
Rachmanino
committed
9b048a4
support threadblock swizzle annotation for cluster launch
Rachmanino
committed
2fc3064
enhance cuda arch restriction for cluster template functions
Rachmanino
committed
8b27b40
Change return type in block_rank_in_cluster function from uint32 to int32
Rachmanino
committed
ae83639
fix threadblock swizzle for cluster
Rachmanino
committed
d45865a
introduce tcgen05.fence::{before, after}_thread_sync
Rachmanino
committed
d2d7111
update lower_blackwell_2sm pass
Rachmanino
committed
20f01be
update b_continuity calculation in GemmTCGEN5 for compatiblity with 2cta
Rachmanino
committed
ebb8b76
use m/n per_cta for offset calculation in 2cta tcgen5 and introduce a config to disable 2cta lowering
Rachmanino
committed
150fb9b
draft refactor of tcgen05mma
Rachmanino
committed
6ddea2a
successfully generated correct 2sm code
Rachmanino
committed
ec9cef9
lint
Rachmanino
committed
dcc1bae
fix bug caused by rebase
Rachmanino
committed
d09fc92
upd 2sm persistent kernel
Rachmanino
committed
a76f221
reorgnize tcgen5mma examples, fix cross-wave bugs, optimize to 1670T
Rachmanino
committed
28617da
lint
Rachmanino
committed
d229491
general support for 2cta in GetTCGEN5MMAMeta
Rachmanino
committed
dbc9388
drop the change of make_tcgen05mma_swizzled_layout
Rachmanino
committed
a01e62d
support 2cta for more dtypes and gemm_ts
Rachmanino
committed
5ddea83
add check for cluster_dims when lowering 2cta tcgen5
Rachmanino
committed
d832d14
refactor 2sm conditional lowering
Rachmanino
committed
97beb2c
remove unnecessary syncthreads
Rachmanino
committed
9572c64
support Layout B and add maint test
Rachmanino
committed
15e3048
lint
Rachmanino
committed
d901fed
Restore support for cp.async.mbarrier and cleaning up unused code
Rachmanino
committed
605d9dc

Commits on Mar 20, 2026

Merge branch 'main' into wt/2sm
Rachmanino
committed
2ec5fe8
Merge branch 'main' into wt/2sm
Rachmanino
committed
4131990
fix tma injection pass change
Rachmanino
committed
33ac055

Commits on Mar 23, 2026

typo
Rachmanino
committed
1caa424
Merge branch 'main' into wt/2sm
Rachmanino
committed
6c09a53
support tma_load_2sm lowering after merge main
Rachmanino
committed
d174892
refactor to expose use_2cta option in `T.tcgen05_gemm` and remove 2cta pass config
Rachmanino
committed
6e698b7
lint
Rachmanino
committed
3795413
fix cutedsl codegen for threadblock swizzle
Rachmanino
committed
d42abc8
lint
Rachmanino
committed
e7e03b6
upd maint script for tcgen05
Rachmanino
committed
28f4224