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

v4.2 release. (#2587)

Browse Source 
* Fix default cluster callback values to 1 to avoid profiler failure when these values are not set in command line.

* v4.2 release.
This commit is contained in:
Junkai-Wu
2025-08-22 18:11:24 -04:00
committed by GitHub
parent 11cad1f67b
commit a49a78ffef
351 changed files with 28182 additions and 2032 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/atom/copy_traits_sm100.hpp
+30
View File
@@ -43,6 +43,36 @@
namespace cute
{
template <>
struct Copy_Traits<SM100_LOAD_256bit_CACHE_NOALLOCATION>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,_256>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,_256>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
};
template <>
struct Copy_Traits<SM100_STORE_256bit_CACHE_NOALLOCATION>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,_256>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,_256>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
};
template <>
struct Copy_Traits<SM100_U8x8_LDSM_T>
{
include/cute/atom/copy_traits_sm100_tma.hpp
+14
View File
@@ -135,6 +135,13 @@ struct Copy_Traits<SM100_TMA_2SM_LOAD_OP, NumBitsPerTMA>
uint64_t*, // smem mbarrier
uint64_t // cache hint
> const opargs_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
//////////////////////////////////////////////////////////////////////////////
@@ -223,6 +230,13 @@ struct Copy_Traits<SM100_TMA_2SM_LOAD_MULTICAST_OP, NumBitsPerTMA>
uint16_t, // multicast mask
uint64_t // cache hint
> const opargs_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
////////////////////////////////////
include/cute/atom/copy_traits_sm75.hpp
-1
View File
@@ -156,5 +156,4 @@ struct Copy_Traits<SM75_U32x1_MOVM_T>
// Reference map from (thr,val) to bit
using RefLayout = DstLayout;
};
} // end namespace cute
include/cute/atom/copy_traits_sm90_tma.hpp
+36 -3
View File
@@ -156,6 +156,13 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
copy_unpack(Copy_Traits const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst) = delete;
// Construct with updated TMA descriptor only (no barrier change)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
with(TmaDescriptor const* new_tma_desc) const {
return {*new_tma_desc, aux_params_};
}
};
// The executable SM90_TMA_LOAD with tma_desc and tma_mbar
@@ -181,6 +188,13 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc, uint64_t* mbar, uint64_t cache)
: opargs_(desc, mbar, cache) {}
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
// The prefetch for SM90_TMA_LOAD with tma_desc
@@ -199,10 +213,22 @@ struct Copy_Traits<SM90_TMA_LOAD::PREFETCH, NumBitsPerTMA, Args...>
tuple<TmaDescriptor const*> const opargs_;
// Construct with any other Traits' TMA Desc
template <class... CopyArgs>
template <class OtherTraits>
CUTE_HOST_DEVICE
Copy_Traits(Copy_Traits<CopyArgs...> const& traits)
: opargs_({&traits.tma_desc_}) {}
Copy_Traits(OtherTraits const& traits)
: opargs_({traits.get_tma_descriptor()}) {}
// Construct directly with a TMA descriptor pointer
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc)
: opargs_({desc}) {}
// Build a new Prefetch traits with a different TMA descriptor pointer
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD::PREFETCH, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc) const {
return {new_tma_desc};
}
template <class TS, class SLayout,
class TD, class DLayout>
@@ -312,6 +338,13 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc, uint64_t* mbar, uint16_t mask, uint64_t hint)
: opargs_(desc, mbar, mask, hint) {}
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
//////////////////////////////////////////////////////////////////////////////
include/cute/atom/mma_traits_sm100.hpp
+156 -6
View File
@@ -2639,10 +2639,10 @@ struct MMA_Traits<SM100_MMA_F8F6F4_SS, a_type, b_type, c_type,
using ValTypeC = c_type;
static_assert(cute::sizeof_bits_v<a_type> <= 8 && cute::sizeof_bits_v<b_type> <= 8, "SM100_MMA_F8F6F4_SS supports types with leq 8bit types");
static_assert(M == 64 || M == 128, "SM100_MMA_F8F6F4_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_F8F6F4_SS N-mode size should be a multiple of 8 between 8 and 256 for M=64,\
or a multiple of 16 between 16 and 256 for M=128.");
static_assert(((b_major == UMMA::Major::K) && ((N % 8 == 0) && (8 <= N) && (N <= 256))) ||
((b_major == UMMA::Major::MN) && ((N % 16 == 0) && (16 <= N) && (N <= 256))),
"SM100_MMA_F8F6F4_SS N-mode size should be a multiple of 8 between 8 and 256 when B is K major. \
SM100_MMA_F8F6F4_SS N-mode size should be a multiple of 16 between 16 and 256 when B is MN major.");
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type>;
@@ -3051,14 +3051,16 @@ struct MMA_Traits<SM100_MMA_F8F6F4_2x1SM_SS, a_type, b_type, c_type,
cute::integral_constant<UMMA::ScaleIn, a_neg>,
cute::integral_constant<UMMA::ScaleIn, b_neg>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
static_assert(cute::sizeof_bits_v<a_type> <= 8 && cute::sizeof_bits_v<b_type> <= 8, "SM100_MMA_F8F6F4_2x1SM_SS supports types with leq 8bit types");
static_assert(M == 128 || M == 256, "SM100_MMA_F8F6F4_2x1SM_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 32 between 32 and 256.");
static_assert(((b_major == UMMA::Major::K) && ((N % 16 == 0) && (16 <= N) && (N <= 256))) ||
((b_major == UMMA::Major::MN) && ((N % 32 == 0) && (32 <= N) && (N <= 256))),
"SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 16 between 16 and 256 when B is K major. \
SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 32 between 32 and 256 when B is MN major.");
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
@@ -3879,4 +3881,152 @@ struct MMA_Traits<SM100_1x2x1_F32F32F32F32>
using CLayout = Layout<Shape<_1,_2>>;
};
namespace SM103 {
// Common mma_unpack for all MMA_Ops in cute::SM103
template <class MMA_Op,
class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr
void
mma_unpack(MMA_Traits<MMA_Op> const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& zA,
Tensor<TB, BLayout> const& zB,
Tensor<TC, CLayout> const& C)
{
auto [A, next_A, SFA] = unzip_tensor(zA);
auto [B, next_B, SFB] = unzip_tensor(zB);
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_rmem<TA>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint64_t desc_a = A[0];
uint64_t desc_next_a = next_A[0];
uint64_t desc_b = B[0];
uint64_t desc_next_b = next_B[0];
auto desc_a_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_a);
auto desc_next_a_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_next_a);
desc_a_temp.lbo_mode_ = 1;
desc_a_temp.leading_byte_offset_ = desc_next_a_temp.start_address_;
auto desc_b_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_b);
auto desc_next_b_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_next_b);
desc_b_temp.lbo_mode_ = 1;
desc_b_temp.leading_byte_offset_ = desc_next_b_temp.start_address_;
uint32_t tmem_c = raw_pointer_cast(D.data());
UMMA::InstrDescriptorBlockScaled instr_desc = traits.idesc_;
instr_desc.k_size_ = 1;
auto tsfa_addr = raw_pointer_cast(SFA.data());
auto tsfb_addr = raw_pointer_cast(SFB.data());
uint64_t idesc = UMMA::make_runtime_instr_desc_block_scaled<>(instr_desc, tsfa_addr, tsfb_addr);
// print("a: "); print(A); print("\n");
// print("b: "); print(B); print("\n");
MMA_Op::fma(reinterpret_cast<uint64_t &>(desc_a_temp), reinterpret_cast<uint64_t &>(desc_b_temp), tmem_c, uint32_t(traits.accumulate_), idesc, tsfa_addr, tsfb_addr);
}
} // end namespace SM103
template <class a_type, class b_type, class c_type, class sf_type,
int M, int N, int VS, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg>
struct MMA_Traits<SM103::SM103_MXF4_ULTRA_SS_VS<a_type, b_type, c_type, sf_type,
M, N, VS, a_major, b_major,
a_neg, b_neg>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using ValTypeSFA = sf_type;
using ValTypeSFB = sf_type;
// Logical shape-K is always 256bits, transform to units of elements
constexpr static int K = 96;
constexpr static int SFVecSize = VS;
static_assert(a_major == UMMA::Major::K && b_major == UMMA::Major::K, "This MMA does not support transpose");
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type>;
using FrgTypeSFA = UMMA::tmem_sf_frg<sf_type, SFVecSize, 1, true>;
using FrgTypeSFB = UMMA::tmem_sf_frg<sf_type, SFVecSize, 1, false>;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_1>;
using ALayout = Layout<Shape <_1,Shape <Int<M>,Int<K>>>,
Stride<_0,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape <_1,Shape <Int<N>,Int<K>>>,
Stride<_0,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape <_1,Shape <Int<M>,Int<N>>>,
Stride<_0,Stride< _1,Int<M>>>>;
using MMA_ScaleFactor = SM100_MMA_MXF4_SS<a_type, b_type, c_type, sf_type,
M, (round_up(N, 128)), VS, a_major, b_major,
a_neg, b_neg>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptorBlockScaled idesc_ = UMMA::make_instr_desc_block_scaled<
a_type, b_type, c_type, sf_type, M, N, a_major, b_major, a_neg, b_neg>();
};
template <class a_type, class b_type, class c_type, class sf_type,
int M, int N, int VS, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg>
struct MMA_Traits<SM103::SM103_MXF4_ULTRA_2x1SM_SS_VS<a_type, b_type, c_type, sf_type,
M, N, VS, a_major, b_major,
a_neg, b_neg>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using ValTypeSFA = sf_type;
using ValTypeSFB = sf_type;
// Logical shape-K is always 256bits, transform to units of elements
constexpr static int K = 96;
constexpr static int SFVecSize = VS;
static_assert(a_major == UMMA::Major::K && b_major == UMMA::Major::K, "This MMA does not support transpose");
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_2sm<c_type>;
constexpr static UMMA::TmemAllocMode TmemAlloc = M == 128 ?
UMMA::TmemAllocMode::ScaleFactorDuplicated2by2 : UMMA::TmemAllocMode::ScaleFactorDuplicated4by1;
using FrgTypeSFA = UMMA::tmem_sf_frg<sf_type, SFVecSize, 2, true, TmemAlloc>;
using FrgTypeSFB = UMMA::tmem_sf_frg<sf_type, SFVecSize, 2, false, TmemAlloc>;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_2>;
using ALayout = Layout<Shape < _2,Shape <Int<M/2>,Int<K>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape < _2,Shape <Int<N/2>,Int<K>>>,
Stride<Int<N/2>,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape < _2,Shape <Int<M/2>,Int<N>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
using MMA_ScaleFactor = SM100_MMA_MXF4_SS<a_type, b_type, c_type, sf_type,
(M/2 > 64 ? M/2 : M), (round_up(N, 128)), VS, a_major, b_major,
a_neg, b_neg>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptorBlockScaled idesc_ = UMMA::make_instr_desc_block_scaled<
a_type, b_type, c_type, sf_type, M, N, a_major, b_major, a_neg, b_neg>();
};
} // end namespace cute
include/cute/atom/mma_traits_sm89.hpp
+1 -2
View File
@@ -67,7 +67,7 @@ struct MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> {
};
template <>
struct MMA_Traits<SM89_16x8x32_F32E4M3E5M2F32_TN>
struct MMA_Traits<SM89_16x8x32_F32E4M3E5M2F32_TN>
: MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> {
using ValTypeD = float;
using ValTypeA = float_e4m3_t;
@@ -129,5 +129,4 @@ struct MMA_Traits<SM89_16x8x32_F16E5M2E4M3F16_TN>
using ValTypeC = cutlass::half_t;
};
} // end namespace cute
include/cute/atom/partitioner.hpp
+2 -1
View File
@@ -31,8 +31,9 @@
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include CUDA_STD_HEADER(type_traits)
#else
#include <type_traits>
#endif