CUTLASS 3.8 Release (#2059)

* CUTLASS 3.8 Release

* update

* Update README.md

* Revert "Update README.md"

This reverts commit b353e36fe83e0815f99b44e46c0c95494c44726b.

* update

* update

---------

Co-authored-by: Haicheng Wu <57973641+hwu36@users.noreply.github.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
mihir-awatramani
2025-01-24 23:44:06 -08:00
committed by GitHub
parent 9eb01fa0b0
commit 389e493055
290 changed files with 91223 additions and 292 deletions

View File

@@ -83,6 +83,14 @@ class DataType(enum.Enum):
s64 = enum_auto()
e4m3 = enum_auto()
e5m2 = enum_auto()
f8 = enum_auto()
f6 = enum_auto()
f4 = enum_auto()
e3m2 = enum_auto()
e2m3 = enum_auto()
e2m1 = enum_auto()
ue8m0 = enum_auto()
ue4m3 = enum_auto()
f16 = enum_auto()
bf16 = enum_auto()
f32 = enum_auto()
@@ -117,6 +125,9 @@ ShortDataTypeNames = {
DataType.f64: 'd',
DataType.cf32: 'c',
DataType.cf64: 'z',
DataType.f8: 'f8',
DataType.f6: 'f6',
DataType.f4: 'f4',
}
#
@@ -137,6 +148,14 @@ DataTypeNames = {
DataType.s64: "s64",
DataType.e4m3: 'e4m3',
DataType.e5m2: 'e5m2',
DataType.f8: 'f8',
DataType.f6: 'f6',
DataType.f4: 'f4',
DataType.e2m3: 'e2m3',
DataType.e3m2: 'e3m2',
DataType.e2m1: 'e2m1',
DataType.ue8m0: 'ue8m0',
DataType.ue4m3: 'ue4m3',
DataType.f16: "f16",
DataType.bf16: "bf16",
DataType.f32: "f32",
@@ -178,6 +197,14 @@ DataTypeTag = {
DataType.s64: "int64_t",
DataType.e4m3: 'cutlass::float_e4m3_t',
DataType.e5m2: 'cutlass::float_e5m2_t',
DataType.f8: 'cutlass::type_erased_dynamic_float8_t',
DataType.f6: 'cutlass::type_erased_dynamic_float6_t',
DataType.f4: 'cutlass::type_erased_dynamic_float4_t',
DataType.e2m3: 'cutlass::float_e2m3_t',
DataType.e3m2: 'cutlass::float_e3m2_t',
DataType.e2m1: 'cutlass::float_e2m1_t',
DataType.ue8m0: 'cutlass::float_ue8m0_t',
DataType.ue4m3: 'cutlass::float_ue4m3_t',
DataType.f16: "cutlass::half_t",
DataType.bf16: "cutlass::bfloat16_t",
DataType.f32: "float",
@@ -219,6 +246,14 @@ DataTypeSize = {
DataType.s64: 64,
DataType.e4m3: 8,
DataType.e5m2: 8,
DataType.f8: 8,
DataType.f6: 6,
DataType.f4: 4,
DataType.e2m3: 6,
DataType.e3m2: 6,
DataType.e2m1: 4,
DataType.ue8m0: 8,
DataType.ue4m3: 8,
DataType.f16: 16,
DataType.bf16: 16,
DataType.f32: 32,
@@ -447,6 +482,22 @@ class KernelScheduleType(enum.Enum):
TmaWarpSpecializedCooperativeFP8FastAccum = enum_auto()
TmaWarpSpecializedPingpongFP8FastAccum = enum_auto()
ImplicitTmaWarpSpecializedSm90 = enum_auto()
TmaWarpSpecialized1SmSm100 = enum_auto()
TmaWarpSpecialized2SmSm100 = enum_auto()
BlockScaledTmaWarpSpecialized1SmSm100 = enum_auto()
BlockScaledTmaWarpSpecialized2SmSm100 = enum_auto()
Mxf8f6f4TmaWarpSpecialized1SmSm100 = enum_auto()
Mxf8f6f4TmaWarpSpecialized2SmSm100 = enum_auto()
Mxf4TmaWarpSpecialized1SmSm100 = enum_auto()
Mxf4TmaWarpSpecialized2SmSm100 = enum_auto()
Nvf4TmaWarpSpecialized1SmSm100 = enum_auto()
Nvf4TmaWarpSpecialized2SmSm100 = enum_auto()
#
KernelScheduleTag = {
KernelScheduleType.ScheduleAuto: 'cutlass::gemm::collective::KernelScheduleAuto',
@@ -462,6 +513,22 @@ KernelScheduleTag = {
KernelScheduleType.TmaWarpSpecializedCooperativeFP8FastAccum: 'cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum',
KernelScheduleType.TmaWarpSpecializedPingpongFP8FastAccum: 'cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum',
KernelScheduleType.ImplicitTmaWarpSpecializedSm90: 'cutlass::conv::KernelImplicitTmaWarpSpecializedSm90',
KernelScheduleType.TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmSm100',
KernelScheduleType.TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmSm100',
KernelScheduleType.BlockScaledTmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledSm100',
KernelScheduleType.BlockScaledTmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledSm100',
KernelScheduleType.Mxf8f6f4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmMxf8f6f4Sm100',
KernelScheduleType.Mxf8f6f4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmMxf8f6f4Sm100',
KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmMxf4Sm100',
KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmMxf4Sm100',
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100',
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100',
}
#
@@ -479,6 +546,22 @@ KernelScheduleSuffixes = {
KernelScheduleType.TmaWarpSpecializedCooperativeFP8FastAccum: '_warpspecialized_cooperative_fp8_fastaccum',
KernelScheduleType.TmaWarpSpecializedPingpongFP8FastAccum: '_warpspecialized_pingpong_fp8_fastaccum',
KernelScheduleType.ImplicitTmaWarpSpecializedSm90: '_warpspecialized',
KernelScheduleType.TmaWarpSpecialized1SmSm100: '_1sm',
KernelScheduleType.TmaWarpSpecialized2SmSm100: '_2sm',
KernelScheduleType.BlockScaledTmaWarpSpecialized1SmSm100: '_1sm',
KernelScheduleType.BlockScaledTmaWarpSpecialized2SmSm100: '_2sm',
KernelScheduleType.Mxf8f6f4TmaWarpSpecialized1SmSm100: '_q_1sm',
KernelScheduleType.Mxf8f6f4TmaWarpSpecialized2SmSm100: '_q_2sm',
KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100: '_o_vs32_1sm',
KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100: '_o_vs32_2sm',
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: '_o_vs16_1sm',
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100: '_o_vs16_2sm',
}
class EpilogueScheduleType(enum.Enum):
@@ -487,6 +570,9 @@ class EpilogueScheduleType(enum.Enum):
NoSmemWarpSpecialized = enum_auto()
TmaWarpSpecialized = enum_auto()
TmaWarpSpecializedCooperative = enum_auto()
TmaWarpSpecialized1Sm = enum_auto()
TmaWarpSpecialized2Sm = enum_auto()
#
EpilogueScheduleTag = {
EpilogueScheduleType.ScheduleAuto: 'cutlass::epilogue::collective::EpilogueScheduleAuto',
@@ -494,6 +580,8 @@ EpilogueScheduleTag = {
EpilogueScheduleType.NoSmemWarpSpecialized: 'cutlass::epilogue::NoSmemWarpSpecialized',
EpilogueScheduleType.TmaWarpSpecialized: 'cutlass::epilogue::TmaWarpSpecialized',
EpilogueScheduleType.TmaWarpSpecializedCooperative: 'cutlass::epilogue::TmaWarpSpecializedCooperative',
EpilogueScheduleType.TmaWarpSpecialized1Sm: 'cutlass::epilogue::TmaWarpSpecialized1Sm',
EpilogueScheduleType.TmaWarpSpecialized2Sm: 'cutlass::epilogue::TmaWarpSpecialized2Sm',
}
#
@@ -503,13 +591,18 @@ EpilogueScheduleSuffixes = {
EpilogueScheduleType.NoSmemWarpSpecialized: '_epi_nosmem',
EpilogueScheduleType.TmaWarpSpecialized: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecializedCooperative: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized1Sm: '',
EpilogueScheduleType.TmaWarpSpecialized2Sm: '_epi_tma',
}
class EpilogueFunctor3x(enum.Enum):
LinearCombination = enum_auto()
LinearCombinationBlockScaleFactor = enum_auto()
#
EpilogueFunctor3xTag = {
EpilogueFunctor3x.LinearCombination: 'cutlass::epilogue::fusion::LinearCombination',
EpilogueFunctor3x.LinearCombinationBlockScaleFactor: 'cutlass::epilogue::fusion::LinCombBlockScaleFactor',
}
class TileSchedulerType(enum.Enum):
@@ -595,12 +688,15 @@ class OpcodeClass(enum.Enum):
TensorOp = enum_auto()
WmmaTensorOp = enum_auto()
SparseTensorOp = enum_auto()
BlockScaledTensorOp = enum_auto()
OpcodeClassNames = {
OpcodeClass.Simt: 'simt',
OpcodeClass.TensorOp: 'tensorop',
OpcodeClass.WmmaTensorOp: 'wmma_tensorop',
OpcodeClass.SparseTensorOp: 'sptensorop',
OpcodeClass.BlockScaledTensorOp: 'bstensorop'
}
OpcodeClassTag = {
@@ -608,6 +704,7 @@ OpcodeClassTag = {
OpcodeClass.TensorOp: 'cutlass::arch::OpClassTensorOp',
OpcodeClass.WmmaTensorOp: 'cutlass::arch::OpClassWmmaTensorOp',
OpcodeClass.SparseTensorOp: 'cutlass::arch::OpClassSparseTensorOp',
OpcodeClass.BlockScaledTensorOp: 'cutlass::arch::OpClassBlockScaledTensorOp'
}
###################################################################################################
@@ -688,6 +785,8 @@ class GemmKind(enum.Enum):
PlanarComplex = enum_auto()
PlanarComplexArray = enum_auto()
Grouped = enum_auto()
BlockScaledUniversal3x = enum_auto()
#
GemmKindNames = {
GemmKind.Gemm: "gemm",
@@ -698,6 +797,7 @@ GemmKindNames = {
GemmKind.PlanarComplex: "gemm_planar_complex",
GemmKind.PlanarComplexArray: "gemm_planar_complex_array",
GemmKind.Grouped: "gemm_grouped",
GemmKind.BlockScaledUniversal3x: "gemm_block_scaled"
}
#
@@ -871,6 +971,8 @@ GroupModeNames = {
GroupMode.Depthwise: 'depthwise',
}
DynamicClusterShape = [0, 0, 1]
###################################################################################################
#
@@ -879,6 +981,7 @@ class MathInstruction:
instruction_shape, \
element_a, element_b, element_accumulator, \
opcode_class, math_operation = MathOperation.multiply_add \
, element_scale_factor = None
):
self.instruction_shape = instruction_shape
@@ -887,6 +990,8 @@ class MathInstruction:
self.element_accumulator = element_accumulator
self.opcode_class = opcode_class
self.math_operation = math_operation
self.element_scale_factor = element_scale_factor
#
class TileDescription: