leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

add pytest support for tutorial gemm (#2826)

Browse Source 
* add pytest support for tutorial gemm

* add license
This commit is contained in:
Linfeng Zheng
2025-12-05 21:45:01 +08:00
committed by GitHub
parent 7252a2d17e
commit f6402fcd5e
7 changed files with 829 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py
View File

@@ -141,15 +141,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(0)
# (MMA, MMA_M, MMA_K)
# (MMA, MMA_M, MMA_K, RestK)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K)
# (MMA, MMA_N, MMA_K, RestK)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K)
# (MMA, MMA_M, MMA_K, STAGE)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K)
# (MMA, MMA_N, MMA_K, STAGE)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])

8
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py
View File

@@ -174,15 +174,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0])
# (MMA, MMA_M, MMA_K)
# (MMA, MMA_M, MMA_K, RestK)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K)
# (MMA, MMA_N, MMA_K, RestK)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K)
# (MMA, MMA_M, MMA_K, STAGE)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K)
# (MMA, MMA_N, MMA_K, STAGE)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])