leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 26 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 15:42:40 -04:00
committed by GitHub
co-authored by yuzhai Haicheng Wu
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
tools/library/include/cutlass/library/descriptions.h
+97 -1
View File
@@ -313,10 +313,16 @@ struct BlockScaleDescription {
TensorDescription SFD;
/// Describes the input ScaleFactor VectorSize
int SFVecSize;
int SFMVecSize;
int SFNVecSize;
int SFKVecSize;
/// Describes the Output ScaleFactor VectorSize
int EpilogueSFVecSize;
/// Describes the underlying kind of scaling:
/// Tensor Core supported (BlockScaled) or manual scaling (Blockwise)
OperationKind kind;
};
struct GroupedGemmDescription : public OperationDescription {
@@ -418,6 +424,96 @@ struct BlockScaledGemmDescription : public OperationDescription {
transform_B(transform_B) {}
};
/// Description of all GEMM computations
struct BlockwiseGemmDescription : public OperationDescription {
/// Indicates the kind of GEMM performed
GemmKind gemm_kind;
/// Describes the A operand
TensorDescription A;
/// Describes the B operand
TensorDescription B;
/// Describes the source matrix
TensorDescription C;
/// Describes the destination matrix
TensorDescription D;
/// Describes the SFA operand
TensorDescription SFA;
/// Describes the SFB operand
TensorDescription SFB;
/// Describes the data type of the scalars passed to the epilogue
NumericTypeID element_epilogue;
/// Describes the structure of parallel reductions
SplitKMode split_k_mode;
/// Transformation on A operand
ComplexTransform transform_A;
/// Transformation on B operand
ComplexTransform transform_B;
/// Describes the input ScaleFactor VectorSize
int SFMVecSize;
int SFNVecSize;
int SFKVecSize;
//
// Methods
//
BlockwiseGemmDescription(
GemmKind gemm_kind = GemmKind::kGemm,
TensorDescription const& A = TensorDescription(),
TensorDescription const& B = TensorDescription(),
TensorDescription const& C = TensorDescription(),
TensorDescription const& D = TensorDescription(),
NumericTypeID element_epilogue = NumericTypeID::kInvalid,
SplitKMode split_k_mode = SplitKMode::kNone,
ComplexTransform transform_A = ComplexTransform::kNone,
ComplexTransform transform_B = ComplexTransform::kNone
):
gemm_kind(gemm_kind),
A(A),
B(B),
C(C),
D(D),
element_epilogue(element_epilogue),
split_k_mode(split_k_mode),
transform_A(transform_A),
transform_B(transform_B) {}
BlockwiseGemmDescription(
OperationDescription op_desc,
GemmKind gemm_kind,
TensorDescription const& A,
TensorDescription const& B,
TensorDescription const& C,
TensorDescription const& D,
NumericTypeID element_epilogue,
SplitKMode split_k_mode,
ComplexTransform transform_A,
ComplexTransform transform_B
):
OperationDescription(op_desc),
gemm_kind(gemm_kind),
A(A),
B(B),
C(C),
D(D),
element_epilogue(element_epilogue),
split_k_mode(split_k_mode),
transform_A(transform_A),
transform_B(transform_B) {}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Description for structured sparse GEMMs.
tools/library/include/cutlass/library/library.h
+65
View File
@@ -121,6 +121,13 @@ public:
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const = 0;
// Set arguments that should only be set once before verifying or profiling the kernel.
// This should encompass any expensive operations that don't vary from run to run
// (e.g., max_active_clusters).
virtual Status initialize_with_arguments(void* arguments_ptr) const {
return Status::kSuccess;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -389,6 +396,56 @@ struct BlockScaledGemmArguments {
bool use_pdl{false};
};
/// Blockwise GEMM
//
// OperationKind: kBlockwiseGemm
// GemmKind: Universal
struct BlockwiseGemmArguments {
// NOTE: these are replicated for 3.0 interfaces
gemm::GemmCoord problem_size{};
gemm::GemmCoord cluster_shape{};
gemm::GemmCoord cluster_shape_fallback{};
int batch_count{1};
void const *A{nullptr};
void const *B{nullptr};
void const *SFA{nullptr};
void const *SFB{nullptr};
void const *C{nullptr};
void *D{nullptr};
void const *alpha{nullptr};
void const *beta{nullptr};
ScalarPointerMode pointer_mode{};
// NOTE: these are replicated for 3.0 interfaces
int64_t lda{0};
int64_t ldb{0};
int64_t ldc{0};
int64_t ldd{0};
int64_t batch_stride_A{0};
int64_t batch_stride_B{0};
int64_t batch_stride_C{0};
int64_t batch_stride_D{0};
int sf_m_vec_size{0};
int sf_n_vec_size{0};
int sf_k_vec_size{0};
// Needed for some 3.x kernels
int sm_count{0};
library::RasterOrder raster_order{};
int swizzle_size{1};
int split_k_slices{1};
library::RuntimeDatatype runtime_input_datatype_a{library::RuntimeDatatype::kStatic};
library::RuntimeDatatype runtime_input_datatype_b{library::RuntimeDatatype::kStatic};
bool use_pdl{false};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -521,6 +578,8 @@ struct GemmGroupedArguments {
// these should really be in the configuration but staying consistent with GEMM
int sm_count{0};
int max_active_clusters{0};
// The user is responsible for allocating storage for problem sizes.
// Since GemmGroupedArguments is used by both the 2.x and 3.x APIs, we
// unfortunately need to have both options in this struct, and the
@@ -536,6 +595,12 @@ struct GroupedGemmBlockScaledArguments : GemmGroupedArguments {
void* norm_constant{nullptr};
};
struct GroupedGemmBlockwiseArguments : GemmGroupedArguments {
void* SFA{nullptr};
void* SFB{nullptr};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// OperationKind: kSparseGemm
tools/library/include/cutlass/library/operation_table.h
+180
View File
@@ -427,6 +427,183 @@ using BlockScaledGemmOperationFunctionalMap = std::unordered_map<
BlockScaledGemmFunctionalKeyHasher
>;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
// Data Structures for Blockwise Gemm Functional Maps
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Tuple uniquely identifying Gemm functional behavior
struct BlockwiseGemmFunctionalKey {
Provider provider;
GemmKind gemm_kind;
OperationKind kind;
NumericTypeID element_compute;
NumericTypeID element_scalar;
NumericTypeID element_A;
LayoutTypeID layout_A;
NumericTypeID element_SFA;
NumericTypeID element_B;
LayoutTypeID layout_B;
NumericTypeID element_SFB;
NumericTypeID element_C;
LayoutTypeID layout_C;
NumericTypeID element_D;
LayoutTypeID layout_D;
int SFMVecSize;
int SFNVecSize;
int SFKVecSize;
//
// Methods
//
inline
BlockwiseGemmFunctionalKey(
Provider provider,
GemmKind gemm_kind = GemmKind::kGemm,
OperationKind kind = OperationKind::kBlockwiseGemm,
NumericTypeID element_compute = NumericTypeID::kF32,
NumericTypeID element_scalar = NumericTypeID::kF32,
NumericTypeID element_A = NumericTypeID::kF16,
LayoutTypeID layout_A = LayoutTypeID::kColumnMajor,
NumericTypeID element_SFA = NumericTypeID::kF16,
NumericTypeID element_B = NumericTypeID::kF16,
LayoutTypeID layout_B = LayoutTypeID::kColumnMajor,
NumericTypeID element_SFB = NumericTypeID::kF16,
NumericTypeID element_C = NumericTypeID::kF16,
LayoutTypeID layout_C = LayoutTypeID::kColumnMajor,
NumericTypeID element_D = NumericTypeID::kF16,
LayoutTypeID layout_D = LayoutTypeID::kColumnMajor,
int sfm_vec_size = 32,
int sfn_vec_size = 32,
int sfk_vec_size = 32
):
provider(provider),
gemm_kind(gemm_kind),
kind(kind),
element_compute(element_compute),
element_scalar(element_scalar),
element_A(element_A),
layout_A(layout_A),
element_SFA(element_SFA),
element_B(element_B),
layout_B(layout_B),
element_SFB(element_SFB),
element_C(element_C),
layout_C(layout_C),
element_D(element_D),
layout_D(layout_D),
SFMVecSize(sfm_vec_size),
SFNVecSize(sfn_vec_size),
SFKVecSize(sfk_vec_size)
{ }
inline
bool operator==(BlockwiseGemmFunctionalKey const &rhs) const {
return
(provider == rhs.provider) &&
(gemm_kind == rhs.gemm_kind) &&
(kind == rhs.kind) &&
(element_compute == rhs.element_compute) &&
(element_scalar == rhs.element_scalar) &&
(element_A == rhs.element_A) &&
(layout_A == rhs.layout_A) &&
(element_SFA == rhs.element_SFA) &&
(element_B == rhs.element_B) &&
(layout_B == rhs.layout_B) &&
(element_SFB == rhs.element_SFB) &&
(element_C == rhs.element_C) &&
(layout_C == rhs.layout_C) &&
(element_D == rhs.element_D) &&
(layout_D == rhs.layout_D) &&
(SFMVecSize == rhs.SFMVecSize) &&
(SFNVecSize == rhs.SFNVecSize) &&
(SFKVecSize == rhs.SFKVecSize);
}
inline
bool operator!=(BlockwiseGemmFunctionalKey const &rhs) const {
return !(*this == rhs);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
inline
std::ostream & operator<<(std::ostream &out, cutlass::library::BlockwiseGemmFunctionalKey const &k) {
out << "{\n"
<< " provider: " << to_string(k.provider) << "\n"
<< " gemm_kind: " << to_string(k.gemm_kind) << "\n"
<< " kind: " << to_string(k.kind) << "\n"
<< " element_compute: " << to_string(k.element_compute) << "\n"
<< " element_scalar: " << to_string(k.element_scalar) << "\n"
<< " element_A: " << to_string(k.element_A) << "\n"
<< " layout_A: " << to_string(k.layout_A) << "\n"
<< " element_SFA: " << to_string(k.element_SFA) << "\n"
<< " element_B: " << to_string(k.element_B) << "\n"
<< " layout_B: " << to_string(k.layout_B) << "\n"
<< " element_SFB: " << to_string(k.element_SFB) << "\n"
<< " element_C: " << to_string(k.element_C) << "\n"
<< " layout_C: " << to_string(k.layout_C) << "\n"
<< " element_D: " << to_string(k.element_D) << "\n"
<< " layout_D: " << to_string(k.layout_D) << "\n"
<< " SFMVecSize: " << k.SFMVecSize << "\n"
<< " SFNVecSize: " << k.SFNVecSize << "\n"
<< " SFKVecSize: " << k.SFKVecSize << "\n"
<< "}";
return out;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Hash function for BlockwiseGemmFunctionalKeyHasher
struct BlockwiseGemmFunctionalKeyHasher {
using IntHash = std::hash<int>;
inline
static size_t rotl(size_t key, int shl) {
return (key << shl) | (key >> (sizeof(key)*8u - static_cast<size_t>(shl)));
}
inline
size_t operator()(BlockwiseGemmFunctionalKey const &key) const {
IntHash hash;
return
rotl(hash(int(key.provider)), 1) ^
rotl(hash(int(key.gemm_kind)), 2) ^
rotl(hash(int(key.kind)), 3) ^
rotl(hash(int(key.element_compute)), 4) ^
rotl(hash(int(key.element_scalar)), 5) ^
rotl(hash(int(key.element_A)), 6) ^
rotl(hash(int(key.layout_A)), 7) ^
rotl(hash(int(key.element_SFA)), 8) ^
rotl(hash(int(key.element_B)), 9) ^
rotl(hash(int(key.layout_B)), 10) ^
rotl(hash(int(key.element_SFB)), 11) ^
rotl(hash(int(key.element_C)), 12) ^
rotl(hash(int(key.layout_C)), 13) ^
rotl(hash(int(key.element_D)), 14) ^
rotl(hash(int(key.layout_D)), 15) ^
rotl(hash(int(key.SFMVecSize)), 16) ^
rotl(hash(int(key.SFNVecSize)), 17) ^
rotl(hash(int(key.SFKVecSize)), 18)
;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Maps a GemmFunctionalKey onto a vector of Operation * objects expected to be of kind kGemm
using BlockwiseGemmOperationFunctionalMap = std::unordered_map<
BlockwiseGemmFunctionalKey,
GemmOperationVectorMap,
BlockwiseGemmFunctionalKeyHasher
>;
/////////////////////////////////////////////////////////////////////////////////////////////////
// Data Structures for Conv Functional Maps
@@ -697,6 +874,9 @@ public:
// provider (kCUTLASS, kReferenceHost, kReferenceDevice)
BlockScaledGemmOperationFunctionalMap block_scaled_gemm_operations;
// provider (kCUTLASS, kReferenceHost, kReferenceDevice)
BlockwiseGemmOperationFunctionalMap blockwise_gemm_operations;
/// Map of all operations of type kConv2d
// provider (kCUTLASS, kReferenceHost, kReferenceDevice)
ConvOperationFunctionalMap conv2d_operations;
tools/library/include/cutlass/library/types.h
+1
View File
@@ -143,6 +143,7 @@ enum class Provider {
enum class OperationKind {
kGemm,
kBlockScaledGemm,
kBlockwiseGemm,
kRankK,
kRank2K,
kTrmm,