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

v4.3 tag release update. (#2789)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-20 20:49:44 -05:00
committed by GitHub
parent 406e078b29
commit 8cd5bef43a
225 changed files with 23229 additions and 2813 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/cutlass_library/gemm_operation.py
+22 -2
View File
@@ -77,6 +77,8 @@ class GemmOperation:
GemmKind.BlockwiseUniversal3x,
GemmKind.GroupedBlockwiseUniversal3x,
GemmKind.BlockScaledSparseUniversal3x,
GemmKind.MoeGroupedUniversal3x,
GemmKind.BlockScaledMoeGroupedUniversal3x,
}
self.is_3x = gemm_kind in kinds_3x
self.prefix = "3x" if self.is_3x else ""
@@ -927,8 +929,14 @@ ${compile_guard_end}
gemm_shape_type = "cute::Shape<int,int,int,int>"
grouped_gemm_shape_type = "cute::Shape<int,int,int>"
grouped_gemm_shape_type = "cutlass::gemm::GroupProblemShape<" + grouped_gemm_shape_type + ">"
return gemm_shape_type if not is_grouped(operation.gemm_kind) else grouped_gemm_shape_type
moe_gemm_shape_type = "cute::Shape<int,int,int>"
moe_gemm_shape_type = "cutlass::gemm::MoEProblemShape<" + moe_gemm_shape_type + ">"
if is_moe(operation.gemm_kind):
return moe_gemm_shape_type
elif is_grouped(operation.gemm_kind):
return grouped_gemm_shape_type
else:
return gemm_shape_type
def emit(self, operation):
_LOGGER.debug("*** EmitGemmConfigurationLibrary::emit(operation)")
@@ -943,6 +951,7 @@ ${compile_guard_end}
instruction_shape = operation.tile_description.math_instruction.instruction_shape
cluster_m = operation.tile_description.cluster_shape[0]
cluster_n = operation.tile_description.cluster_shape[1]
cta_m = tile_shape[0] // cluster_m if cluster_m > 0 else tile_shape[0]
cta_n = tile_shape[1] // cluster_n if cluster_n > 0 else tile_shape[1]
tile_shape_m, tile_shape_n, tile_shape_k = operation.get_collective_tile_shape()
@@ -1023,6 +1032,13 @@ ${compile_guard_end}
element_a = f'cute::tuple<{str(element_a)},{str(DataTypeTag[operation.ScaleFactorA])}>'
element_b = f'cute::tuple<{str(element_b)},{str(DataTypeTag[operation.ScaleFactorB])}>'
if is_moe(operation.gemm_kind):
if DataTypeSize[operation.A.element] == 4 and operation.ScaleFactorA == DataType.ue4m3:
element_a = f"cutlass::nv_float4_t<{DataTypeTag[operation.A.element]}>"
if DataTypeSize[operation.B.element] == 4 and operation.ScaleFactorB == DataType.ue4m3:
element_b = f"cutlass::nv_float4_t<{DataTypeTag[operation.B.element] }>"
alignment_c = get_tma_alignment(operation.C.element) \
if is_tma_epilogue(operation.epilogue_schedule) and opcode_class_epi != OpcodeClass.Simt \
else operation.C.alignment
@@ -1480,6 +1496,8 @@ class EmitGemmConfigurationLibrary:
GemmKind.BlockwiseUniversal3x: EmitGemmUniversal3xInstance,
GemmKind.GroupedBlockwiseUniversal3x: EmitGemmUniversal3xInstance,
GemmKind.BlockScaledSparseUniversal3x: EmitGemmUniversal3xInstance,
GemmKind.MoeGroupedUniversal3x: EmitGemmUniversal3xInstance,
GemmKind.BlockScaledMoeGroupedUniversal3x: EmitGemmUniversal3xInstance,
}
self.gemm_kind_wrappers = {
@@ -1497,6 +1515,8 @@ class EmitGemmConfigurationLibrary:
GemmKind.BlockwiseUniversal3x: 'BlockwiseGemmUniversal3xOperation',
GemmKind.GroupedBlockwiseUniversal3x: 'GroupedBlockwiseGemmUniversal3xOperation',
GemmKind.BlockScaledSparseUniversal3x: 'BlockScaledSparseGemmUniversal3xOperation',
GemmKind.MoeGroupedUniversal3x: 'MoeGroupedGemmUniversal3xOperation',
GemmKind.BlockScaledMoeGroupedUniversal3x: 'BlockScaledMoeGroupedGemmUniversal3xOperation',
}
self.wmma_guard_start = "#if defined(CUTLASS_ARCH_WMMA_SM${sm_number}_ENABLED)"
python/cutlass_library/generator.py
+701 -14
View File
@@ -5841,7 +5841,24 @@ def GenerateSM90_TensorOp_fp8_WGMMA_gemm_with_blockwise(manifest, cuda_version,
level=instantiation_level)
tile_descriptions = list()
tile_descriptions.append(
TileDescription(
threadblock_shape=[
256,
128,
128
],
stages=0,
warp_count=[4, 1, 1],
math_instruction=MathInstruction(
[128, 128, 32],
DataType.e5m2, DataType.e4m3, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add),
min_compute=90,
max_compute=90,
cluster_shape=[1,2,1],
explicit_vector_sizes=[1, 128, 128]))
for desc in tile_descriptions_:
desc.explicit_vector_sizes = [1, desc.tile_shape[1], desc.tile_shape[2]]
tile_descriptions.append(copy.deepcopy(desc))
@@ -7043,6 +7060,104 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed,
[[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
def GenerateSM100_TensorOp_16b_UMMA_alignx_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999)
# layouts for ABC and their alignments. C alignment will be set later based on output type
layouts = [
[[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]],
[[LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]],
[[LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]],
[[LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]],
[[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]],
[[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]],
[[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]],
[[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]],
]
thor_sm = ThorSMRenumbering(cuda_version)
math_instructions_1sm, _ = generate_16b_math_instructions_sm100(instantiation_level)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
grouped = is_grouped(gemm_kind)
if grouped:
return
cluster_shapes_1sm= [[1,1,1]]
tile_schedulers = [
TileSchedulerType.Default
]
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in cluster_shapes_1sm:
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_1sm[0],
math_inst.instruction_shape[1] * multiplier_1sm[1],
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
data_types = [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : math_inst.element_accumulator,
"d_type" : math_inst.element_accumulator,
"acc_type" : math_inst.element_accumulator,
"epi_type" : math_inst.element_accumulator,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : math_inst.element_accumulator,
"acc_type" : math_inst.element_accumulator,
"epi_type" : math_inst.element_accumulator,
},
]
kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100
epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
[[kernel_schedule, epi_schedule]],
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
# for mixed precision kernels, also generate kernels that write output matrix in the A/B format
# Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation)
if math_inst.element_a != math_inst.element_accumulator:
data_types_mixed = [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : math_inst.element_a,
"d_type" : math_inst.element_a,
"acc_type" : math_inst.element_accumulator,
"epi_type" : math_inst.element_accumulator,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : math_inst.element_a,
"acc_type" : math_inst.element_accumulator,
"epi_type" : math_inst.element_accumulator,
},
]
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed,
[[kernel_schedule, epi_schedule]],
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
@@ -7198,16 +7313,19 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
# Set alignment d based on Destination format.
for layout in layouts:
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
for data_type in data_types:
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
( data_type["d_type"] == DataType.e5m2 ):
continue
kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized1SmSm100, grouped)
epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, epi_schedule]],
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
for tile_description in tile_descriptions:
for layout in layouts:
for data_type in data_types:
if layout[1][0] == LayoutType.RowMajor and tile_description.math_instruction.instruction_shape[1] % 16 != 0:
continue
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
( data_type["d_type"] == DataType.e5m2 ):
continue
kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized1SmSm100, grouped)
epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)
CreateGemmUniversal3xOperator(manifest, [layout], [tile_description], data_type,
[[kernel_schedule, epi_schedule]],
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
# 2xSM MMA kernels
@@ -7341,6 +7459,158 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
instantiation_level = manifest.get_instantiation_level(pruned_level=591 , default_level=591 , exhaustive_level=9999)
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], # TN Layout
[[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], # TN Layout
]
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
grouped = is_grouped(gemm_kind)
math_instructions_1sm, _ = generate_fp8_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped)
cluster_shapes_1sm = [[1,1,1]]
tile_schedulers = [
TileSchedulerType.Default
]
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in cluster_shapes_1sm:
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_1sm[0],
math_inst.instruction_shape[1] * multiplier_1sm[1],
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
data_types = [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.f16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.bf16,
"d_type" : DataType.bf16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.bf16,
"d_type" : DataType.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.bf16,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f32,
"d_type" : DataType.f32,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.f16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.bf16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.f32,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
}
]
for data_type in data_types:
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
( data_type["d_type"] == DataType.e5m2 ):
continue
kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100
epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, epi_schedule]],
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.BlockwiseUniversal3x):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
@@ -7467,7 +7737,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version,
continue
is_runtime_datatype_a = is_runtime_datatype(data_type["a_type"])
is_runtime_datatype_b = is_runtime_datatype(data_type["d_type"])
is_runtime_datatype_b = is_runtime_datatype(data_type["b_type"])
# A/B datatypes should be both static or dynamic
if (is_runtime_datatype_a != is_runtime_datatype_b):
@@ -8585,6 +8855,418 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifes
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
)
def GenerateSM100_TensorOp_16b_UMMA_moe_gemm(manifest, cuda_version, gemm_kind=GemmKind.MoeGroupedUniversal3x):
# SM100 MOE GEMM
if not CudaToolkitVersionSatisfies(cuda_version, 13, 0):
return
instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999)
# layouts for ABC and their alignments. C alignment will be set later based on output type
layouts = [
[[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 0]],
]
thor_sm = ThorSMRenumbering(cuda_version)
math_instructions_1sm, math_instructions_2sm = generate_16b_math_instructions_sm100(instantiation_level)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
cluster_shapes= [[1,1,1]]
tile_schedulers = [
TileSchedulerType.Default
]
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in cluster_shapes:
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_1sm[0],
math_inst.instruction_shape[1] * multiplier_1sm[1],
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
data_types = [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : math_inst.element_accumulator,
"d_type" : math_inst.element_b,
"acc_type" : math_inst.element_accumulator,
"epi_type" : math_inst.element_accumulator,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : math_inst.element_b,
"acc_type" : math_inst.element_accumulator,
"epi_type" : math_inst.element_accumulator,
},
]
# Set alignment d based on Destination format.
for layout in layouts:
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
[[kernel_schedule, epi_schedule]],
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
def GenerateSM100_TensorOp_fp8_UMMA_moe_gemm(manifest, cuda_version, gemm_kind=GemmKind.MoeGroupedUniversal3x):
# SM100 MOE GEMM
if not CudaToolkitVersionSatisfies(cuda_version, 13, 0):
return
instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999)
# layouts for ABC and their alignments. C alignment will be set later based on output type
layouts = [
[[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]],
]
thor_sm = ThorSMRenumbering(cuda_version)
math_instructions_1sm, math_instructions_2sm = generate_fp8_math_instructions_sm100(instantiation_level)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
# only support 1x1x1 cluster shape
cluster_shapes= [[1,1,1]]
epi_type = DataType.f32
tile_schedulers = [
TileSchedulerType.Default
]
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in cluster_shapes:
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_1sm[0],
math_inst.instruction_shape[1] * multiplier_1sm[1],
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
data_types = [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.f16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.bf16,
"d_type" : DataType.bf16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.bf16,
"d_type" : DataType.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
},
]
# Set alignment d based on Destination format.
for layout in layouts:
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm
for data_type in data_types:
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
( data_type["d_type"] == DataType.e5m2 ):
continue
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, epi_schedule]],
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
def GenerateSM100_TensorOp_mixed_8bits_UMMA_moe_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledMoeGroupedUniversal3x):
# SM100 moe GEMM with mixed F4/F6/F8 inputs + block scale
if not CudaToolkitVersionSatisfies(cuda_version, 13, 0):
return
instantiation_level = manifest.get_instantiation_level(pruned_level=590, default_level=590, exhaustive_level=9999)
grouped = is_grouped(gemm_kind)
layouts = [
[[LayoutType.RowMajor, 0], [LayoutType.ColumnMajor, 0], [LayoutType.ColumnMajor, 0]],
]
math_instructions_1sm, math_instructions_2sm = generate_mxf8f6f4_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped)
cluster_shapes_1sm = [[1,1,1]]
acc_types = [ DataType.f32 ]
def tile_schedulers(sfdtype):
# Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void,
# the epilogue is the traditional linear combination, for which we already have tests with stream-K.
if sfdtype["type"] == DataType.void or grouped:
return [TileSchedulerType.Default]
else:
return [TileSchedulerType.Default, TileSchedulerType.StreamK]
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8)
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp
if DataTypeSize[math_inst.element_a] != DataTypeSize[math_inst.element_b] or DataTypeSize[math_inst.element_a] == 6:
continue
tile_descriptions = []
for cluster_shape in cluster_shapes_1sm:
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_1sm[0],
math_inst.instruction_shape[1] * multiplier_1sm[1],
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
data_types = [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.bf16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e3m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
}]
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm
for data_type in data_types:
# Set alignment d based on Destination format.
for layout in layouts:
layout[0][1] = 128 // DataTypeSize[data_type["a_type"]]
layout[1][1] = 128 // DataTypeSize[data_type["b_type"]]
layout[2][1] = 128 // DataTypeSize[data_type["d_type"]]
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
schedules=[[kernel_schedule, epi_schedule]], tile_schedulers=[TileSchedulerType.Default], gemm_kind=gemm_kind)
def GenerateSM100_TensorOp_fp4_UMMA_MoE_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledMoeGroupedUniversal3x):
# SM100 MoE GEMM with F4 + block scale
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999)
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]],
]
shapes_1sm = [
(128, 64, 64), (128, 128, 64), (128, 192, 64), (128, 256, 64)
]
math_instructions_1sm = []
for shape in shapes_1sm:
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.e2m1, DataType.e2m1, DataType.f32,
OpcodeClass.BlockScaledTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
cluster_shapes_1sm = [[1,1,1]]
acc_types = [ DataType.f32 ] # Accumulator is always 32 bits for block scaled MMA instructions
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8)
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp
tile_descriptions = []
for cluster_shape in cluster_shapes_1sm:
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_1sm[0],
math_inst.instruction_shape[1] * multiplier_1sm[1],
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
assert math_inst.instruction_shape[2] * 4 == 256
data_types = [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.bf16,
"d_type" : DataType.bf16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.e2m1,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e5m2,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : DataType.e2m1,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e2m1,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor}
},
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.f16,
"d_type" : DataType.e2m1,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor}
}
]
# Set alignment d based on Destination format.
for layout in layouts:
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
for layout in layouts:
for data_type in data_types:
if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor):
data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout.
if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor):
continue
# E2M1 x E2M1, vector size 32, E8
# E2M1 x E2M1, vector size 16, UE4M3
isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100
epi_schedule = EpilogueScheduleType.ScheduleAuto
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, [[kernel_schedule, epi_schedule]]
, tile_schedulers=[TileSchedulerType.Default], gemm_kind=gemm_kind
)
def GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledUniversal3x):
# SM100 MMA with F4 + block scale
@@ -11047,17 +11729,22 @@ def GenerateSM100(manifest, cuda_version):
# Dense Gemm
#
GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version)
GenerateSM100_TensorOp_16b_UMMA_alignx_gemm(manifest, cuda_version)
GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version)
if not bool(set(manifest.compute_capabilities_feature_set).intersection(arch_family_cc)):
GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version)
GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version)
GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version)
# grouped GEMM
GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x)
GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x)
# MOE grouped GEMM
GenerateSM100_TensorOp_16b_UMMA_moe_gemm(manifest, cuda_version)
GenerateSM100_TensorOp_fp8_UMMA_moe_gemm(manifest, cuda_version)
GenerateSM100_TensorOp_mixed_8bits_UMMA_moe_gemm_with_block_scaled(manifest, cuda_version)
GenerateSM100_TensorOp_fp4_UMMA_MoE_gemm_with_block_scaled(manifest, cuda_version)
# StreamK is included in regular generation
GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version)
python/cutlass_library/library.py
+22 -3
View File
@@ -322,7 +322,8 @@ def is_complex(data_type):
return False
def is_block_scaled(gemm_kind):
return gemm_kind in (GemmKind.BlockScaledUniversal3x, GemmKind.GroupedBlockScaledUniversal3x, GemmKind.BlockScaledSparseUniversal3x)
return gemm_kind in (GemmKind.BlockScaledUniversal3x, GemmKind.GroupedBlockScaledUniversal3x, GemmKind.BlockScaledSparseUniversal3x,
GemmKind.BlockScaledMoeGroupedUniversal3x)
def is_blockwise(gemm_kind):
return gemm_kind in (GemmKind.BlockwiseUniversal3x, GemmKind.GroupedBlockwiseUniversal3x)
@@ -331,6 +332,8 @@ def is_grouped(gemm_kind):
return gemm_kind in (GemmKind.GroupedUniversal3x,
GemmKind.GroupedBlockScaledUniversal3x, GemmKind.GroupedBlockwiseUniversal3x)
def is_moe(gemm_kind):
return gemm_kind in (GemmKind.MoeGroupedUniversal3x, GemmKind.BlockScaledMoeGroupedUniversal3x)
#
def get_complex_from_real(real_type):
for r, c in RealComplexBijection:
@@ -513,6 +516,8 @@ class KernelScheduleType(enum.Enum):
TmaWarpSpecialized1SmSm100 = enum_auto()
TmaWarpSpecialized2SmSm100 = enum_auto()
WarpSpecialized1SmSm100 = enum_auto()
ImplicitTmaWarpSpecialized1SmSm100 = enum_auto()
ImplicitTmaWarpSpecialized2SmSm100 = enum_auto()
@@ -528,6 +533,9 @@ class KernelScheduleType(enum.Enum):
PtrArrayMxf8f6f4TmaWarpSpecialized1SmSm100 = enum_auto()
PtrArrayMxf8f6f4TmaWarpSpecialized2SmSm100 = enum_auto()
MixedTmaCpAsyncWarpSpecialized1SmSm100 = enum_auto()
MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100 = enum_auto()
SparseTmaWarpSpecialized1SmSm100 = enum_auto()
SparseTmaWarpSpecialized2SmSm100 = enum_auto()
@@ -621,6 +629,7 @@ KernelScheduleTag = {
KernelScheduleType.TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmSm100',
KernelScheduleType.TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmSm100',
KernelScheduleType.WarpSpecialized1SmSm100: 'cutlass::gemm::KernelWarpSpecialized1SmSm100',
KernelScheduleType.ImplicitTmaWarpSpecialized1SmSm100: 'cutlass::conv::KernelImplicitTmaWarpSpecialized1SmSm100',
KernelScheduleType.ImplicitTmaWarpSpecialized2SmSm100: 'cutlass::conv::KernelImplicitTmaWarpSpecialized2SmSm100',
@@ -641,7 +650,8 @@ KernelScheduleTag = {
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockwise1SmSm100',
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockwise2SmSm100',
KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100: 'cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmSm100',
KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100: 'cutlass::gemm::KernelMixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100',
KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmMxf4Sm100',
KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmMxf4Sm100',
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100',
@@ -738,6 +748,7 @@ KernelScheduleSuffixes = {
KernelScheduleType.TmaWarpSpecialized1SmSm100: '_1sm',
KernelScheduleType.TmaWarpSpecialized2SmSm100: '_2sm',
KernelScheduleType.WarpSpecialized1SmSm100: '_cpasync_1sm',
KernelScheduleType.ImplicitTmaWarpSpecialized1SmSm100: '_1sm',
KernelScheduleType.ImplicitTmaWarpSpecialized2SmSm100: '_2sm',
@@ -758,6 +769,9 @@ KernelScheduleSuffixes = {
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized1SmSm100: '_1sm',
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized2SmSm100: '_2sm',
KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100: '_mixed_cpasync_1sm',
KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100: '_mixed_cpasync_1sm',
KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100: '_o_vs32_1sm',
KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100: '_o_vs32_2sm',
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: '_o_vs16_1sm',
@@ -1208,6 +1222,9 @@ class GemmKind(enum.Enum):
BlockwiseUniversal3x = enum_auto()
GroupedBlockwiseUniversal3x = enum_auto()
BlockScaledSparseUniversal3x = enum_auto()
MoeGroupedUniversal3x = enum_auto()
BlockScaledMoeGroupedUniversal3x = enum_auto()
#
GemmKindNames = {
@@ -1224,7 +1241,9 @@ GemmKindNames = {
GemmKind.GroupedBlockScaledUniversal3x: "gemm_grouped",
GemmKind.BlockwiseUniversal3x: "gemm",
GemmKind.GroupedBlockwiseUniversal3x: "gemm_grouped",
GemmKind.BlockScaledSparseUniversal3x: "spgemm"
GemmKind.BlockScaledSparseUniversal3x: "spgemm",
GemmKind.MoeGroupedUniversal3x: "moe_gemm",
GemmKind.BlockScaledMoeGroupedUniversal3x: "moe_gemm",
}
#