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

cutlass 3.9 update (#2255)

Browse Source 
* cutlass 3.9 update

* rebase

* fixes out of shared memory for blockwise Blackwell

* doc format

* fix issue 2253

* disable host ref by default

* fix sm120 smem capacity

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-04-24 12:42:40 -07:00
committed by GitHub
parent 8e345c5c5b
commit 331a1f5b3f
143 changed files with 18089 additions and 5935 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
examples/cute/tutorial/blackwell/01_mma_sm100.cu
View File

@@ -61,7 +61,8 @@
#include <cute/tensor.hpp> // CuTe tensor implementation
#include <cute/arch/cluster_sm90.hpp> // CuTe functions for querying the details of cluster launched
#include <cute/numeric/integral_constant.hpp> // Compile time in constants such as _1, _256 etc.
#include <cute/algorithm/cooperative_copy.hpp>
#include <cute/algorithm/cooperative_copy.hpp> // Auto vectorized copy operation
#include <cute/arch/tmem_allocator_sm100.hpp> // TMEM allocator for SM100
// Tutorial helpers
#include "example_utils.hpp"
@@ -122,7 +123,9 @@ struct SharedStorage
alignas(128) cute::ArrayEngine<TypeA, cute::cosize_v<ASmemLayout>> A;
alignas(128) cute::ArrayEngine<TypeB, cute::cosize_v<BSmemLayout>> B;
alignas(16) cute::uint64_t mma_barrier; // Barrier to track MMA computation on SMEM
alignas(16) cute::uint64_t mma_barrier; // Barrier to track MMA computation on SMEM
alignas(16) cute::uint32_t tmem_base_ptr; // Base pointer for TMEM allocation
CUTE_DEVICE constexpr auto tensor_sA() { return make_tensor(make_smem_ptr(A.begin()), ASmemLayout{}); }
CUTE_DEVICE constexpr auto tensor_sB() { return make_tensor(make_smem_ptr(B.begin()), BSmemLayout{}); }
@@ -225,6 +228,18 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
// ThrMma's make_fragment_C() creates a TMEM tensor with the appropriate layout for the accumulator.
Tensor tCtAcc = cta_mma.make_fragment_C(tCgC); // (MmaC, NumMma_M, NumMma_N)
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
using TmemAllocator = cute::TMEM::Allocator1Sm;
TmemAllocator tmem_allocator{};
if (elect_one_warp) {
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
}
__syncthreads(); // Wait for all threads until warp0 allocates TMEM
tCtAcc.data() = shared_storage.tmem_base_ptr;
if (thread0()) {
print("tCsA:\t"); print(tCsA); print("\n"); // tCsA: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_A) o ((_128,_16),_1,_4):((_64,_1),_0,_16)
print("tCsB:\t"); print(tCsB); print("\n"); // tCsB: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_B) o ((_256,_16),_1,_4):((_64,_1),_0,_16)
@@ -233,10 +248,8 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
print("tCtAcc:\t"); print(tCtAcc); print("\n"); // tCtAcc: tmem_[32b](TMEM_ADDR) o ((_128,_256),_1,_1):((_65536,_1),_0,_0)
} __syncthreads();
// Barrier Initialization
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
// Barrier Initialization
// Barriers in SMEM initialized by a single thread.
if (elect_one_warp && elect_one_thr) {
cute::initialize_barrier(shared_storage.mma_barrier, /* num_ctas */ 1);
@@ -306,6 +319,15 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
axpby(alpha, tDrAcc, beta, tDrC);
// Store RMEM -> GMEM
copy(tDrC, tDgD);
__syncthreads();
// Release the right to allocate before deallocations so that the next CTA can rasterize
// Then deallocate TMEM
if (elect_one_warp) {
tmem_allocator.release_allocation_lock();
tmem_allocator.free(shared_storage.tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
}
template <class TypeA, class LayoutA,

29
examples/cute/tutorial/blackwell/02_mma_tma_sm100.cu
View File

@@ -61,7 +61,8 @@
#include <cute/tensor.hpp> // CuTe tensor implementation
#include <cute/arch/cluster_sm90.hpp> // CuTe functions for querying the details of cluster launched
#include <cute/numeric/integral_constant.hpp> // Compile time in constants such as _1, _256 etc.
#include <cute/algorithm/cooperative_copy.hpp>
#include <cute/algorithm/cooperative_copy.hpp> // Auto vectorized copy operation
#include <cute/arch/tmem_allocator_sm100.hpp> // TMEM allocator for SM100
// Tutorial helpers
#include "example_utils.hpp"
@@ -124,6 +125,8 @@ struct SharedStorage
alignas(16) cute::uint64_t mma_barrier; // Barrier to track MMA computation on SMEM
alignas(16) cute::uint64_t tma_barrier; // Barrier to track TMA data transfers to SMEM
alignas(16) cute::uint32_t tmem_base_ptr; // Base pointer for TMEM allocation
CUTE_DEVICE constexpr auto tensor_sA() { return make_tensor(make_smem_ptr(A.begin()), ASmemLayout{}); }
CUTE_DEVICE constexpr auto tensor_sB() { return make_tensor(make_smem_ptr(B.begin()), BSmemLayout{}); }
};
@@ -228,6 +231,18 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
// ThrMma's make_fragment_C() creates a TMEM tensor with the appropriate layout for the accumulator.
Tensor tCtAcc = cta_mma.make_fragment_C(tCgC); // (MmaC, NumMma_M, NumMma_N)
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
using TmemAllocator = cute::TMEM::Allocator1Sm;
TmemAllocator tmem_allocator{};
if (elect_one_warp) {
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
}
__syncthreads(); // Wait for all threads until warp0 allocates TMEM
tCtAcc.data() = shared_storage.tmem_base_ptr;
if (thread0()) {
print("tCsA:\t"); print(tCsA); print("\n"); // tCsA: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_A) o ((_128,_16),_1,_4):((_64,_1),_0,_16)
print("tCsB:\t"); print(tCsB); print("\n"); // tCsB: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_B) o ((_256,_16),_1,_4):((_64,_1),_0,_16)
@@ -269,9 +284,6 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
} __syncthreads();
// Barrier Initialization
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
// Barriers in SMEM initialized by a single thread.
if (elect_one_warp && elect_one_thr) {
cute::initialize_barrier(shared_storage.mma_barrier, /* num_ctas */ 1);
@@ -346,6 +358,15 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
axpby(alpha, tDrAcc, beta, tDrC);
// Store RMEM -> GMEM
copy(tDrC, tDgD);
__syncthreads();
// Release the right to allocate before deallocations so that the next CTA can rasterize
// Then deallocate TMEM
if (elect_one_warp) {
tmem_allocator.release_allocation_lock();
tmem_allocator.free(shared_storage.tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
}
template <class TypeA, class LayoutA,

30
examples/cute/tutorial/blackwell/03_mma_tma_multicast_sm100.cu
View File

@@ -61,7 +61,8 @@
#include <cute/tensor.hpp> // CuTe tensor implementation
#include <cute/arch/cluster_sm90.hpp> // CuTe functions for querying the details of cluster launched
#include <cute/numeric/integral_constant.hpp> // Compile time in constants such as _1, _256 etc.
#include <cute/algorithm/cooperative_copy.hpp>
#include <cute/algorithm/cooperative_copy.hpp> // Auto vectorized copy operation
#include <cute/arch/tmem_allocator_sm100.hpp> // TMEM allocator for SM100
// Tutorial helpers
#include "example_utils.hpp"
@@ -129,6 +130,8 @@ struct SharedStorage
alignas(16) cute::uint64_t mma_barrier; // Barrier to track MMA computation on SMEM
alignas(16) cute::uint64_t tma_barrier; // Barrier to track TMA data transfers to SMEM
alignas(16) cute::uint32_t tmem_base_ptr; // Base pointer for TMEM allocation
CUTE_DEVICE constexpr auto tensor_sA() { return make_tensor(make_smem_ptr(A.begin()), ASmemLayout{}); }
CUTE_DEVICE constexpr auto tensor_sB() { return make_tensor(make_smem_ptr(B.begin()), BSmemLayout{}); }
};
@@ -231,6 +234,18 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
// ThrMma's make_fragment_C() creates a TMEM tensor with the appropriate layout for the accumulator.
Tensor tCtAcc = cta_mma.make_fragment_C(tCgC); // (MmaC, NumMma_M, NumMma_N)
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
using TmemAllocator = cute::TMEM::Allocator1Sm;
TmemAllocator tmem_allocator{};
if (elect_one_warp) {
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
}
__syncthreads(); // Wait for all threads until warp0 allocates TMEM
tCtAcc.data() = shared_storage.tmem_base_ptr;
if (thread0()) {
print("tCsA:\t"); print(tCsA); print("\n"); // tCsA: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_A) o ((_128,_16),_1,_4):((_64,_1),_0,_16)
print("tCsB:\t"); print(tCsB); print("\n"); // tCsB: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_B) o ((_256,_16),_1,_4):((_64,_1),_0,_16)
@@ -305,10 +320,6 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
} __syncthreads();
// Barrier Initialization
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
// Barriers in SMEM initialized by a single thread.
if (elect_one_warp && elect_one_thr) {
// The number of CTAs that participates in multicast operation with this CTA (for both A and B matrices)
@@ -385,6 +396,15 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
axpby(alpha, tDrAcc, beta, tDrC);
// Store RMEM -> GMEM
copy(tDrC, tDgD);
__syncthreads();
// Release the right to allocate before deallocations so that the next CTA can rasterize
// Then deallocate TMEM
if (elect_one_warp) {
tmem_allocator.release_allocation_lock();
tmem_allocator.free(shared_storage.tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
}
template <class TypeA, class LayoutA,

31
examples/cute/tutorial/blackwell/04_mma_tma_2sm_sm100.cu
View File

@@ -61,7 +61,8 @@
#include <cute/tensor.hpp> // CuTe tensor implementation
#include <cute/arch/cluster_sm90.hpp> // CuTe functions for querying the details of cluster launched
#include <cute/numeric/integral_constant.hpp> // Compile time in constants such as _1, _256 etc.
#include <cute/algorithm/cooperative_copy.hpp>
#include <cute/algorithm/cooperative_copy.hpp> // Auto vectorized copy operation
#include <cute/arch/tmem_allocator_sm100.hpp> // TMEM allocator for SM100
// Tutorial helpers
#include "example_utils.hpp"
@@ -132,6 +133,8 @@ struct SharedStorage
alignas(16) cute::uint64_t mma_barrier; // Barrier to track MMA computation on SMEM
alignas(16) cute::uint64_t tma_barrier; // Barrier to track TMA data transfers to SMEM
alignas(16) cute::uint32_t tmem_base_ptr; // Base pointer for TMEM allocation
CUTE_DEVICE constexpr auto tensor_sA() { return make_tensor(make_smem_ptr(A.begin()), ASmemLayout{}); }
CUTE_DEVICE constexpr auto tensor_sB() { return make_tensor(make_smem_ptr(B.begin()), BSmemLayout{}); }
};
@@ -234,6 +237,18 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
// ThrMma's make_fragment_C() creates a TMEM tensor with the appropriate layout for the accumulator.
Tensor tCtAcc = cta_mma.make_fragment_C(tCgC); // (MmaC, NumMma_M, NumMma_N)
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
using TmemAllocator = cute::TMEM::Allocator2Sm;
TmemAllocator tmem_allocator{};
if (elect_one_warp) {
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
}
__syncthreads(); // Wait for all threads until warp0 allocates TMEM
tCtAcc.data() = shared_storage.tmem_base_ptr;
if (thread0()) {
print("tCsA:\t"); print(tCsA); print("\n"); // tCsA: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_A) o ((_128,_16),_1,_4):((_64,_1),_0,_16)
print("tCsB:\t"); print(tCsB); print("\n"); // tCsB: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_B) o ((_256,_16),_1,_4):((_64,_1),_0,_16)
@@ -262,6 +277,7 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
// Construct the CTA-in-Cluster coordinate for multicasting
auto cta_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(int(cute::block_rank_in_cluster()));
auto elect_one_cta = get<0>(cta_in_cluster_coord_vmnk) == Int<0>{};
// Project the cluster_layout for tma_A along the N-modes
auto [tAgA, tAsA] = tma_partition(tma_atom_A,
@@ -299,10 +315,6 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
} __syncthreads();
// Barrier Initialization
auto elect_one_thr = cute::elect_one_sync();
auto elect_one_warp = (threadIdx.x / 32 == 0);
auto elect_one_cta = get<0>(cta_in_cluster_coord_vmnk) == Int<0>{};
// Barriers in SMEM should be initialized by a single thread.
if (elect_one_warp && elect_one_thr) {
// The number of CTAs that participates in multicast operation with this CTA (for both A and B matrices)
@@ -386,6 +398,15 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
axpby(alpha, tDrAcc, beta, tDrC);
// Store RMEM -> GMEM
copy(tDrC, tDgD);
__syncthreads();
// Release the right to allocate before deallocations so that the next CTA can rasterize
// Then deallocate TMEM
if (elect_one_warp) {
tmem_allocator.release_allocation_lock();
tmem_allocator.free(shared_storage.tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
}
template <class TypeA, class LayoutA,

30
examples/cute/tutorial/blackwell/05_mma_tma_epi_sm100.cu
View File

@@ -61,7 +61,8 @@
#include <cute/tensor.hpp> // CuTe tensor implementation
#include <cute/arch/cluster_sm90.hpp> // CuTe functions for querying the details of cluster launched
#include <cute/numeric/integral_constant.hpp> // Compile time in constants such as _1, _256 etc.
#include <cute/algorithm/cooperative_copy.hpp>
#include <cute/algorithm/cooperative_copy.hpp> // Auto vectorized copy operation
#include <cute/arch/tmem_allocator_sm100.hpp> // TMEM allocator for SM100
// Tutorial helpers
#include "example_utils.hpp"
@@ -140,6 +141,8 @@ struct SharedStorage
alignas(16) cute::uint64_t mma_barrier; // Barrier to track MMA computation on SMEM
alignas(16) cute::uint64_t tma_barrier; // Barrier to track TMA data transfers to SMEM
alignas(16) cute::uint32_t tmem_base_ptr; // Base pointer for TMEM allocation
CUTE_DEVICE constexpr auto tensor_sA() { return make_tensor(make_smem_ptr(tensors.mainloop.A.begin()), ASmemLayout{}); }
CUTE_DEVICE constexpr auto tensor_sB() { return make_tensor(make_smem_ptr(tensors.mainloop.B.begin()), BSmemLayout{}); }
CUTE_DEVICE constexpr auto tensor_sC() { return make_tensor(make_smem_ptr(tensors.C.begin()), CSmemLayout{}); }
@@ -247,6 +250,18 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
// ThrMma's make_fragment_C() creates a TMEM tensor with the appropriate layout for the accumulator.
Tensor tCtAcc = cta_mma.make_fragment_C(tCgC); // (MmaC, NumMma_M, NumMma_N)
uint32_t elect_one_thr = cute::elect_one_sync();
uint32_t elect_one_warp = (threadIdx.x / 32 == 0);
using TmemAllocator = cute::TMEM::Allocator2Sm;
TmemAllocator tmem_allocator{};
if (elect_one_warp) {
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
}
__syncthreads(); // Wait for all threads until warp0 allocates TMEM
tCtAcc.data() = shared_storage.tmem_base_ptr;
if (thread0()) {
print("tCsA:\t"); print(tCsA); print("\n"); // tCsA: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_A) o ((_128,_16),_1,_4):((_64,_1),_0,_16)
print("tCsB:\t"); print(tCsB); print("\n"); // tCsB: Sw<3,4,3>_smem_ptr[16b](SMEM_ADDR_B) o ((_256,_16),_1,_4):((_64,_1),_0,_16)
@@ -275,6 +290,7 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
// Construct the CTA-in-Cluster coordinate for multicasting
auto cta_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(int(cute::block_rank_in_cluster()));
auto elect_one_cta = get<0>(cta_in_cluster_coord_vmnk) == Int<0>{};
// Project the cluster_layout for tma_A along the N-modes
auto [tAgA, tAsA] = tma_partition(tma_atom_A,
@@ -312,10 +328,6 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
} __syncthreads();
// Barrier Initialization
auto elect_one_thr = cute::elect_one_sync();
auto elect_one_warp = (threadIdx.x / 32 == 0);
auto elect_one_cta = get<0>(cta_in_cluster_coord_vmnk) == Int<0>{};
// Barriers in SMEM should be initialized by a single thread.
if (elect_one_warp && elect_one_thr) {
// The number of CTAs that participates in multicast operation with this CTA (for both A and B matrices)
@@ -441,6 +453,14 @@ gemm_device(ATensor mA, // (Gemm_M, Gemm_K)
}
__syncthreads(); // All threads sync with issuing thread
}
__syncthreads();
// Release the right to allocate before deallocations so that the next CTA can rasterize
// Then deallocate TMEM
if (elect_one_warp) {
tmem_allocator.release_allocation_lock();
tmem_allocator.free(shared_storage.tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
}
template <class TypeA, class LayoutA,