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

v4.0 update. (#2371)

Browse Source
This commit is contained in:
Junkai-Wu
2025-06-06 14:39:20 +08:00
committed by GitHub
parent 2e2af190bd
commit 8bdbfca682
254 changed files with 29751 additions and 1980 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/cutlass_library/__init__.py
View File

@@ -35,6 +35,7 @@ import sys
from . import conv2d_operation
from . import conv3d_operation
from . import emit_kernel_listing
from . import gemm_operation
if '-m' not in sys.argv:
@@ -53,7 +54,7 @@ from . import trmm_operation
from .library import *
# Set up `source` to point to the path containing the CUTLASS source.
# Check first if the path cotains a `source` subdirectory -- this will
# Check first if the path contains a `source` subdirectory -- this will
# be the case when the package has been installed via pip. Otherwise,
# default to the root of CUTLASS.
install_source_path = os.path.join(__path__[0], 'source')

87
python/cutlass_library/emit_kernel_listing.py
View File

@@ -253,7 +253,8 @@ def _getInstType(input_precision, accumulate_precision, math_instruction):
return inst
# TODO: Computes FLOps/Bytes for GEMM - revisit for conv
def _computeFlopsPerByte(operation, m, n, k, batch_count=1, beta=0.0):
def _computeFlopsPerByte(operation, m, n, k, batch_count=1, beta=0.0, num_groups=1):
assert not (batch_count > 1 and num_groups > 1)
# TODO: adjust for sparsity
gmem_bytes = (
@@ -269,16 +270,15 @@ def _computeFlopsPerByte(operation, m, n, k, batch_count=1, beta=0.0):
gmem_bytes += (DataTypeSize[operation.C.element] * m // 8) * n
flops += 2 * m * n
gmem_bytes *= batch_count
flops *= batch_count
multiplier = max(batch_count, num_groups)
gmem_bytes *= multiplier
flops *= multiplier
return flops / gmem_bytes
def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
):
profiler_reference_computing = "--verification-providers=device --providers=cutlass"
# beta values for L0 and L1
# TODO: randomize beta values for wider coverage
beta_values = [0.5]
@@ -303,15 +303,9 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
]
sm100_mma_data_type_runtime_dtype = [
'gemm_f4_f4_f32_f32_f32',
'gemm_f6_f6_f32_f32_f32',
'gemm_f8_f8_f32_f32_f32',
]
sm100_mma_data_type_mergeable = [
'gemm_e4m3_e4m3_f32_f32_f32',# mask out one instance for verification
'gemm_e2m1_e2m1_f32_f32_f32',
'gemm_e3m2_e3m2_f32_f32_f32',
'gemm.*f4_f4_f32_f32_f32',
'gemm.*f6_f6_f32_f32_f32',
'gemm.*f8_f8_f32_f32_f32',
]
sm100_mma_cluster_size = [
@@ -327,9 +321,6 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
]
# regex list must be in kernel procedural name order
mergeable_sm100_mma_filter_regex_1sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_data_type_mergeable, sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*1sm.*"
mergeable_sm100_mma_filter_regex_2sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_data_type_mergeable, sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*2sm.*"
sm100_mma_filter_regex_1sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_data_type_general, sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*1sm.*"
sm100_mma_filter_regex_2sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_data_type_general, sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*2sm.*"
@@ -340,25 +331,15 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
# Block Scale Gemm
#
block_scaled_data_type_base = [
block_scaled_data_type = [
# runtime datatypes
'gemm.*ue8m0xf4_ue8m0xf4_f32_f16_e5m2',
'gemm.*ue8m0xf4_ue8m0xf4_f32_f16_e5m2',
'gemm.*ue4m3xf4_ue4m3xf4_f32_f16_e5m2',
'gemm.*ue8m0xf4_ue8m0xf6_f32_f16_e5m2',
'gemm.*ue8m0xf4_ue8m0xf4_f32_f16_ue8m0xe2m1',
'gemm.*ue8m0xf6_ue8m0xf6_f32_f16_ue8m0xe3m2',
]
block_scaled_data_type_mergeable = [
'gemm.*ue8m0xe2m1_ue8m0xe2m1_f32_f16_e5m2',
'gemm.*ue8m0xe2m1_ue8m0xe2m1_f32_f16_e5m2',
'gemm.*ue8m0xe2m1_ue8m0xe2m3_f32_f16_e5m2',
'gemm.*ue8m0xe2m1_ue8m0xe2m1_f32_f16_ue8m0xe2m1',
'gemm.*ue8m0xe2m3_ue8m0xe2m3_f32_f16_ue8m0xe3m2',
]
block_scaled_data_type = block_scaled_data_type_base + block_scaled_data_type_mergeable
block_scaled_cluster_size = [
'4x4x1', '2x1x1',
'0x0x1' # dynamic cluster
@@ -366,27 +347,25 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
block_scaled_layouts = ['tnt']
# regex list must be in kernel procedural name order
mergeable_block_scaled_filter_regex_1sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type_mergeable, block_scaled_cluster_size, block_scaled_layouts]]) + ").*1sm.*"
mergeable_block_scaled_filter_regex_2sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type_mergeable, block_scaled_cluster_size, block_scaled_layouts]]) + ").*2sm.*"
block_scaled_filter_regex_1sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type, block_scaled_cluster_size, block_scaled_layouts]]) + ").*1sm.*"
block_scaled_filter_regex_2sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type, block_scaled_cluster_size, block_scaled_layouts]]) + ").*2sm.*"
if arch == "100a" or arch == "100f":
if arch in ["100a", "100f"]:
kernel_filter = f"({sm100_mma_filter_regex_1sm})|" \
f"({sm100_mma_filter_regex_2sm})|" \
f"({sm100_mma_filter_regex_1sm_runtime})|" \
f"({sm100_mma_filter_regex_2sm_runtime})|" \
f"({block_scaled_filter_regex_1sm})|" \
f"({block_scaled_filter_regex_2sm})"
elif arch == "101a" or arch == "101f":
elif arch in ["101a", "101f",
]:
kernel_filter = f"({sm100_mma_filter_regex_1sm})|" \
f"({sm100_mma_filter_regex_2sm})|" \
f"({sm100_mma_filter_regex_1sm_runtime})|" \
f"({sm100_mma_filter_regex_2sm_runtime})|" \
f"({block_scaled_filter_regex_1sm})|" \
f"({block_scaled_filter_regex_2sm})"
elif arch == "120a" or arch == "120f":
elif arch in ["120a", "120f"]:
# blockscaled sm120_mma kernels
blockscaled_sm120_mma_kernel_cta_tiles = [
@@ -403,18 +382,8 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
else:
error_message = "unsupported arch, only support sm100a, sm100f, sm101a, sm101f, sm120a, sm120f"
raise Exception(error_message)
# Statically encoded kernels are still added to generated_kernels
# but are filtered out from the testing commands to reduce test duration.
# The mergeable_kernel_filter specifies the kernels that are already covered
# by the runtime datatype tests so that we safely mark them off
# without changing the test coverage.
mergeable_kernel_filter = f"({mergeable_sm100_mma_filter_regex_1sm})|" \
f"({mergeable_sm100_mma_filter_regex_2sm})|" \
f"({mergeable_block_scaled_filter_regex_1sm})|" \
f"({mergeable_block_scaled_filter_regex_2sm})"
elif mode == "functional_L1":
elif mode == "functional_L1":
sm100_mma_cluster_size = [
'0x0x1' # dynamic cluster
]
@@ -486,10 +455,6 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
audit_file_params_name = os.path.join(curr_build_dir, f"FK_{mode}_audit_params_SM{arch}_cutlass3x_gemm.csv")
if is_runtime_datatype_enabled:
mergeable_kernel_filter_re = re.compile(mergeable_kernel_filter)
kernel_filter_re = re.compile(kernel_filter)
testcase_counter = 0
kernels_emitted = 0
@@ -517,12 +482,6 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
if 'f16_f16_f16_void_f16' not in kernel_name :
continue
# Filter out the statically encoded tests which are
# covered by runtime datatype tests to avoid repetition.
if is_runtime_datatype_enabled and len(mergeable_kernel_filter_re.findall(kernel_name)) != 0:
continue
kernels_emitted += 1
kernel_name_set.add(kernel_name)
hashed_kernel_name = hash_cutlass_string(kernel_name)
@@ -685,9 +644,18 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
gemm_op = "gemm"
profiler_reference_computing_override = profiler_reference_computing
grouped = is_grouped(manifest.operations_by_name[kernel_name].gemm_kind)
num_groups = 1
if "bstensorop" in kernel_name:
profiler_reference_computing_override = "--mode=trace"
if grouped:
gemm_op = "grouped_gemm"
num_groups = 3 # small to limit test time in host block-scaled reference kernels
batch_count = 1
elif "bstensorop" in kernel_name:
gemm_op = "block_scaled_gemm"
elif is_blockwise(manifest.operations_by_name[kernel_name].gemm_kind):
gemm_op = "blockwise_gemm"
problem_size_category = ['smallK','largeK'][index_k] + '_' + ['beta==0','beta!=0'][bool(beta)]
@@ -704,7 +672,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
'n' : n,
'k' : k,
'beta' : beta,
'flops_per_byte' : _computeFlopsPerByte(operation, m, n, k, batch_count, beta)
'flops_per_byte' : _computeFlopsPerByte(operation, m, n, k, batch_count, beta, num_groups)
},
"runtime_params": {
'ctas_per_mma_instruction' : ctas_per_mma_instruction,
@@ -732,6 +700,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
f" --m={str(m)}" +
f" --n={str(n)}" +
f" --k={str(k)}" +
(f" --num_groups={str(num_groups)}" if grouped else "") +
f" --cluster_m={str(cluster_shape_m)}" +
f" --cluster_n={str(cluster_shape_n)}" +
f" --cluster_k={str(cluster_shape_k)}" +
@@ -739,7 +708,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
f" --cluster_n_fallback={str(cluster_n_fallback)}" +
f" --cluster_k_fallback={str(cluster_k_fallback)}" +
f" --beta={str(beta)}" +
f" --batch_count={str(batch_count)}" +
("" if grouped else f" --batch_count={str(batch_count)}") +
f" --swizzle_size={str(swizzle_size)}" +
f" --verification-required={str(verification_required).lower()}"
] \
@@ -752,7 +721,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
testcase_metadata.append(json.dumps(metadata_dict))
testlist_csv_rows.append(testcase_metadata)
testcase_counter += 1
alpha = 1.0
if dynamic_datatype:

10
python/cutlass_library/gemm_operation.py
View File

@@ -994,6 +994,12 @@ ${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])}>'
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
alignment_d = get_tma_alignment(operation.D.element) \
if is_tma_epilogue(operation.epilogue_schedule) and opcode_class_epi != OpcodeClass.Simt \
else operation.D.alignment
operation_name_str = operation.procedural_name()
layout_a_str = LayoutTag[instance_layout_A]
@@ -1103,8 +1109,8 @@ using {operation_name_str}_LayoutSFB = decltype({operation_name_str}_ScaleConfig
'stages': stage_count_string,
'align_a': str(operation.A.alignment),
'align_b': str(operation.B.alignment),
'align_c': str(operation.C.alignment),
'align_d': str(operation.C.alignment),
'align_c': str(alignment_c),
'align_d': str(alignment_d),
'transform_a': ComplexTransformTag[operation.A.complex_transform],
'transform_b': ComplexTransformTag[operation.B.complex_transform],
'math_operation': MathOperationTag[operation.tile_description.math_instruction.math_operation],

314
python/cutlass_library/generator.py
View File

@@ -112,7 +112,6 @@ def CudaToolkitVersionSatisfies(semantic_ver_string, major, minor, patch = 0):
cuda_version.append(x)
return cuda_version >= [major, minor, patch]
###################################################################################################
###################################################################################################
@@ -6769,8 +6768,9 @@ def GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version):
},
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
math_instructions_1sm = [
# tf32 -> f32
MathInstruction(
@@ -6793,8 +6793,8 @@ def GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version):
cluster_shapes_1sm = [[1,2,1], [1,1,1], [1,4,1], [4,4,1]
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline:
cluster_shapes_1sm = [[1,2,1], [1,1,1], [1,4,1]
, DynamicClusterShape
]
@@ -6847,7 +6847,7 @@ def GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]
, DynamicClusterShape
]
@@ -6887,8 +6887,9 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
[[LayoutType.RowMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.RowMajor, 0]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
grouped = is_grouped(gemm_kind)
math_instructions_1sm = [
@@ -6950,7 +6951,7 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [[1,2,1], [1,1,1], [1,4,1]
, DynamicClusterShape
]
@@ -7108,7 +7109,7 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]
, DynamicClusterShape
]
@@ -7152,7 +7153,7 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2Sm
else:
epi_schedule = EpilogueScheduleType.ScheduleAuto
kernel_schedule = KernelScheduleType.TmaWarpSpecialized2SmSm100 if not grouped else KernelScheduleType.PtrArrayTmaWarpSpecialized2SmSm100
kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized2SmSm100, grouped)
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
[[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
@@ -7201,8 +7202,9 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
[[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
epi_type = DataType.f32
grouped = is_grouped(gemm_kind)
@@ -7270,7 +7272,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [[1,2,1], [2,1,1], [1,1,1], [1,4,1]
, DynamicClusterShape
]
@@ -7398,11 +7400,8 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
( data_type["d_type"] == DataType.e5m2 ):
continue
# don't support runtime data type for grouped yet
if grouped and (data_type["a_type"] == DataType.f8 or data_type["b_type"] == DataType.f8):
continue
kernel_schedule = KernelScheduleType.TmaWarpSpecialized1SmSm100 if not grouped else KernelScheduleType.PtrArrayTmaWarpSpecialized1SmSm100
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm if not grouped else EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm
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)
@@ -7484,7 +7483,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]
, DynamicClusterShape
]
@@ -7607,9 +7606,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
( data_type["d_type"] == DataType.e5m2 ):
continue
# don't support runtime data type for grouped yet
if grouped and (data_type["a_type"] == DataType.f8 or data_type["b_type"] == DataType.f8):
continue
if grouped:
epi_schedule = EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm
@@ -7617,7 +7613,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2Sm
else:
epi_schedule = EpilogueScheduleType.ScheduleAuto
kernel_schedule = KernelScheduleType.TmaWarpSpecialized2SmSm100 if not grouped else KernelScheduleType.PtrArrayTmaWarpSpecialized2SmSm100
kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized2SmSm100, grouped)
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
@@ -7852,9 +7848,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version,
if (is_runtime_datatype_a != is_runtime_datatype_b):
continue
# grouped GEMM does not support runtime data type yet
if grouped and (is_runtime_datatype_a or is_runtime_datatype_b):
continue
kernel_schedule = to_grouped_schedule(KernelScheduleType.BlockwiseTmaWarpSpecialized1SmSm100, grouped)
epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
@@ -7896,8 +7889,9 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
TileSchedulerType.Default, TileSchedulerType.StreamK
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
epi_type = DataType.f32
math_instructions_1sm = []
@@ -7949,7 +7943,7 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [
[2,1,1],
[1,1,1]
@@ -8025,7 +8019,7 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [
[2,1,1]
, DynamicClusterShape
@@ -8131,8 +8125,9 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cud
else:
return [TileSchedulerType.Default, TileSchedulerType.StreamK]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
epi_type = DataType.f32
math_instructions_1sm = []
@@ -8184,7 +8179,7 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cud
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [
[1,1,1],
[2,1,1]
@@ -8264,7 +8259,7 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cud
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [
[2,1,1],
[4,1,1]
@@ -8372,6 +8367,7 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]],
[[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]],
]
instruction_sizes_1sm = [
@@ -8400,8 +8396,9 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
else:
return [TileSchedulerType.Default, TileSchedulerType.StreamK]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
epi_type = DataType.f32
math_instructions_1sm = []
@@ -8416,10 +8413,6 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
if (is_runtime_datatype_a != is_runtime_datatype_b):
continue
# grouped GEMM does not support runtime data type yet
if grouped and (is_runtime_datatype_a or is_runtime_datatype_b):
continue
math_instructions_1sm.append(
MathInstruction(
instr_size,
@@ -8447,10 +8440,6 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
if (is_runtime_datatype_a != is_runtime_datatype_b):
continue
# grouped GEMM does not support runtime data type yet
if grouped and (is_runtime_datatype_a or is_runtime_datatype_b):
continue
math_instructions_2sm.append(
MathInstruction(
instr_size,
@@ -8477,7 +8466,7 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [
[1,1,1],
[2,1,1]
@@ -8575,8 +8564,11 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
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):
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
@@ -8604,7 +8596,7 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [
[2,1,1],
[4,1,1]
@@ -8701,8 +8693,11 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
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):
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
isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1
@@ -8737,8 +8732,9 @@ def GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version):
[[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
epi_type = DataType.f32
@@ -8763,7 +8759,7 @@ def GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [[1,2,1], [2,1,1], [1,1,1], [1,4,1]
, DynamicClusterShape
]
@@ -8867,7 +8863,7 @@ def GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]
, DynamicClusterShape
]
@@ -8952,8 +8948,9 @@ def GenerateSM100_SparseTensorOp_32b_UMMA_gemm(manifest, cuda_version):
[[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
tile_schedulers = [
TileSchedulerType.Default,
]
@@ -9009,7 +9006,7 @@ def GenerateSM100_SparseTensorOp_32b_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
@@ -9040,7 +9037,7 @@ def GenerateSM100_SparseTensorOp_32b_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_2sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2])
@@ -9077,8 +9074,9 @@ def GenerateSM100_SparseTensorOp_16b_UMMA_gemm(manifest, cuda_version):
[[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
tile_schedulers = [
TileSchedulerType.Default,
]
@@ -9134,7 +9132,7 @@ def GenerateSM100_SparseTensorOp_16b_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
@@ -9165,7 +9163,7 @@ def GenerateSM100_SparseTensorOp_16b_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_2sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2])
@@ -9202,8 +9200,9 @@ def GenerateSM100_SparseTensorOp_int8_UMMA_gemm(manifest, cuda_version):
[[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
tile_schedulers = [
TileSchedulerType.Default,
@@ -9259,7 +9258,7 @@ def GenerateSM100_SparseTensorOp_int8_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
@@ -9290,7 +9289,7 @@ def GenerateSM100_SparseTensorOp_int8_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_2sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2])
@@ -9327,8 +9326,9 @@ def GenerateSM100_SparseTensorOp_fp8_UMMA_gemm(manifest, cuda_version):
[[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
tile_schedulers = [
TileSchedulerType.Default,
]
@@ -9389,7 +9389,7 @@ def GenerateSM100_SparseTensorOp_fp8_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
@@ -9424,7 +9424,7 @@ def GenerateSM100_SparseTensorOp_fp8_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_2sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2])
@@ -9465,8 +9465,9 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
[[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
tile_schedulers = [
TileSchedulerType.Default,
]
@@ -9525,7 +9526,7 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
@@ -9587,7 +9588,7 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_2sm:
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
if cluster_shape == [4,4,1] :
continue
multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2])
@@ -9677,8 +9678,9 @@ def GenerateSM100_TensorOp_32b_UMMA_gemm_stream_k(manifest, cuda_version):
}
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
math_instructions_1sm = [
MathInstruction(
[128, 256, 8],
@@ -9692,7 +9694,7 @@ def GenerateSM100_TensorOp_32b_UMMA_gemm_stream_k(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [
[1,2,1], [1,1,1], [1,4,1]
, DynamicClusterShape
@@ -9732,7 +9734,7 @@ def GenerateSM100_TensorOp_32b_UMMA_gemm_stream_k(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [
[2,1,1], [2,2,1], [2,4,1], [4,1,1]
, DynamicClusterShape
@@ -9770,8 +9772,9 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm_stream_k(manifest, cuda_version):
[[LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.RowMajor, 0]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
math_instructions_1sm = [
MathInstruction(
[128, 256, 16],
@@ -9784,7 +9787,7 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm_stream_k(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [
[1,2,1], [1,1,1]
, DynamicClusterShape
@@ -9858,7 +9861,7 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm_stream_k(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [
[2,1,1], [2,2,1], [2,4,1], [4,1,1]
, DynamicClusterShape
@@ -9931,8 +9934,9 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm_stream_k(manifest, cuda_version):
[[LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 0]],
]
thor_sm = 101
min_cc = 100
max_cc = 101
max_cc = thor_sm
epi_type = DataType.f32
math_instructions_1sm = [
@@ -9947,7 +9951,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm_stream_k(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [
[1,2,1], [2,1,1], [1,1,1]
, DynamicClusterShape
@@ -10005,7 +10009,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm_stream_k(manifest, cuda_version):
, DynamicClusterShape
]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [
[2,1,1], [2,2,1], [2,4,1], [4,1,1]
, DynamicClusterShape
@@ -10080,8 +10084,9 @@ def GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version,
if not CudaToolkitVersionSatisfies(cuda_version, 12, 0):
return
thor_sm = 101
minimum_compute_capability = 100
maximum_compute_capability = 101
maximum_compute_capability = thor_sm
spatial_dims = [2, 3]
@@ -10110,7 +10115,7 @@ def GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version,
cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1],[4,4,1]]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1]]
# tile_descriptions is a 2-level list.
@@ -10176,7 +10181,7 @@ def GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version,
data_types_and_instruction_shapes_2sm)
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1], [4,4,1]]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]]
for math_inst, output_type in math_instructions_w_output_2sm:
@@ -10233,8 +10238,9 @@ def GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version,
if not CudaToolkitVersionSatisfies(cuda_version, 12, 0):
return
thor_sm = 101
minimum_compute_capability = 100
maximum_compute_capability = 101
maximum_compute_capability = thor_sm
spatial_dims = [2, 3]
stages = 0 # zero means "deduce the number of stages automatically"
@@ -10258,7 +10264,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version,
data_types_and_instruction_shapes_1sm)
cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1],[4,4,1]]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1]]
for math_inst, output_type in math_instructions_w_output_1sm:
@@ -10323,7 +10329,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version,
data_types_and_instruction_shapes_2sm)
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1], [4,4,1]]
if 101 in manifest.compute_capabilities :
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]]
for math_inst, output_type in math_instructions_w_output_2sm:
@@ -10704,9 +10710,9 @@ def GenerateSM120_Sparse_TensorOp_gemm(manifest, cuda_version):
ab_types_mxf8f6f4 = [
DataType.e2m1,
DataType.e2m3,
#DataType.e2m3,
DataType.e3m2,
DataType.e5m2,
#DataType.e5m2,
DataType.e4m3,
]
@@ -10783,13 +10789,145 @@ def GenerateSM120_Sparse_TensorOp_gemm(manifest, cuda_version):
tile_schedulers = tile_schedulers(kernel_schedule),
gemm_kind = GemmKind.SparseUniversal3x)
def GenerateSM120_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.BlockwiseUniversal3x):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
layouts = [
[[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 16]],
[[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.ColumnMajor, 16]]
]
cooperative_tile_sizes = [
[128, 128, 128]
]
pingpong_tile_sizes = [
[64, 128, 128]
]
def get_tile_sizes(kernel_scheduler):
if kernel_scheduler == KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120:
return pingpong_tile_sizes
return cooperative_tile_sizes
def get_warp_count(kernel_scheduler):
if kernel_scheduler == KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120:
return [2, 2, 1]
return [4, 2, 1]
def get_sf_sizes(tile_size):
sf_sizes = []
for vec_m in [1, 128]:
if tile_size[0] % vec_m > 0:
continue
for vec_n in [1, 128]:
if tile_size[1] % vec_m > 0:
continue
sf_sizes.append(
[vec_m, vec_n, 128]
)
return sf_sizes
cluster_shape = [1,1,1]
acc_types = [ DataType.f32 ]
instruction_sizes = [
[16, 8, 32]
]
def tile_schedulers(kernel_schedule):
return [TileSchedulerType.Default]
min_cc = 120
max_cc = 120
kernel_schedulers = [
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperativeSm120,
KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120
]
ab_types = [
[DataType.e4m3, DataType.e4m3],
[DataType.e4m3, DataType.e5m2]
]
math_instructions = []
for instr_size, ab_type, acc_type in product(instruction_sizes, ab_types, acc_types):
a_type, b_type = ab_type
math_instructions.append(
MathInstruction(
instr_size,
a_type, b_type, acc_type,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
# Create gemm operator for mxf8f6f4
for kernel_schedule in kernel_schedulers:
tile_sizes = get_tile_sizes(kernel_schedule)
warp_count = get_warp_count(kernel_schedule)
for math_inst in math_instructions:
tile_descriptions = []
for tile_size in tile_sizes:
sf_sizes = get_sf_sizes(tile_size)
for sf_size in sf_sizes:
tile_descriptions.append(
TileDescription(tile_size, 0, warp_count, math_inst, min_cc, max_cc, cluster_shape,
explicit_vector_sizes=sf_size)
)
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" : DataType.f32
},
{
"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" : DataType.f32
},
{
"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" : DataType.f32
},
{
"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" : DataType.f32
}
]
for data_type in data_types:
# Set alignment d based on Destination format
for layout in layouts:
layout[2][1] = int(128 // DataTypeSize[data_type["d_type"]])
# Create gemm operator
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, EpilogueScheduleType.ScheduleAuto]],
tile_schedulers = tile_schedulers(kernel_schedule),
gemm_kind = gemm_kind)
def GenerateSM100(manifest, cuda_version):
arch_family_cc = ['100f', '101f']
#
# Dense Gemm
#
architectures = manifest.args.architectures.split(';') if len(args.architectures) else ['50',]
GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version)
GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version)
@@ -10797,7 +10935,7 @@ def GenerateSM100(manifest, cuda_version):
GenerateSM100_TensorOp_16b_UMMA_gemm_stream_k(manifest, cuda_version)
if '100f' not in architectures and '101f' not in architectures:
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)
@@ -10819,7 +10957,7 @@ def GenerateSM100(manifest, cuda_version):
#
GenerateSM100_SparseTensorOp_32b_UMMA_gemm(manifest, cuda_version)
GenerateSM100_SparseTensorOp_16b_UMMA_gemm(manifest, cuda_version)
if '100f' not in architectures and '101f' not in architectures:
if not bool(set(manifest.compute_capabilities_feature_set).intersection(arch_family_cc)):
GenerateSM100_SparseTensorOp_int8_UMMA_gemm(manifest, cuda_version)
GenerateSM100_SparseTensorOp_fp8_UMMA_gemm(manifest, cuda_version)
GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version)
@@ -10849,6 +10987,8 @@ def GenerateSM120(manifest, cuda_version):
# Sparse Gemm
#
GenerateSM120_Sparse_TensorOp_gemm(manifest, cuda_version)
GenerateSM120_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version)
GenerateSM120_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockwiseUniversal3x)
###################################################################################################
@@ -11328,13 +11468,17 @@ if __name__ == "__main__":
GenerateSM80(manifest, args.cuda_version)
GenerateSM89(manifest, args.cuda_version)
GenerateSM90(manifest, args.cuda_version)
blackwell_enabled_arch = any(arch in ["100a", "100f", "101a", "101f", "120a", "120f"] for arch in archs)
blackwell_arch_list = [
"100a", "100f",
"101a", "101f",
"120a", "120f"
]
blackwell_enabled_arch = any(arch in blackwell_arch_list for arch in archs)
if blackwell_enabled_arch:
GenerateSM100(manifest, args.cuda_version)
GenerateSM120(manifest, args.cuda_version)
if 'library' in args.generator_target.split(','):
manifest.emit(GeneratorTarget.Library)

43
python/cutlass_library/library.py
View File

@@ -345,6 +345,15 @@ def get_real_from_complex(complex_type):
return r
return DataType.invalid
# TMA requires an alignment of 128 bits for all data types
def get_tma_alignment(data_type):
if data_type == DataType.void:
return 0
elif DataTypeSize[data_type] == 6:
return 128 # 96B alignment for 16U6 format
else:
return 128 // DataTypeSize[data_type]
#
class ComplexMultiplyOp(enum.Enum):
multiply_add = enum_auto()
@@ -546,6 +555,9 @@ class KernelScheduleType(enum.Enum):
F8f6f4SparseTmaWarpSpecializedCooperativeSm120 = enum_auto()
BlockwiseTmaWarpSpecializedCooperativeSm120 = enum_auto()
BlockwiseTmaWarpSpecializedPingpongSm120 = enum_auto()
KernelScheduleTag = {
KernelScheduleType.ScheduleAuto: 'cutlass::gemm::collective::KernelScheduleAuto',
KernelScheduleType.Multistage: 'cutlass::gemm::KernelMultistage',
@@ -614,7 +626,10 @@ KernelScheduleTag = {
KernelScheduleType.Mxf4TmaWarpSpecializedCooperativeSm120: 'cutlass::gemm::KernelTmaWarpSpecializedMxf4Sm120',
KernelScheduleType.Mxf4TmaWarpSpecializedPingpongSm120: 'cutlass::gemm::KernelTmaWarpSpecializedPingpongMxf4Sm120',
KernelScheduleType.F8f6f4SparseTmaWarpSpecializedCooperativeSm120: 'cutlass::gemm::KernelScheduleSparseF8f6f4Sm120'
KernelScheduleType.F8f6f4SparseTmaWarpSpecializedCooperativeSm120: 'cutlass::gemm::KernelScheduleSparseF8f6f4Sm120',
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperativeSm120: 'cutlass::gemm::KernelTmaWarpSpecializedBlockwiseCooperativeSm120',
KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120: 'cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120',
}
#
@@ -685,7 +700,10 @@ KernelScheduleSuffixes = {
KernelScheduleType.Mxf4TmaWarpSpecializedCooperativeSm120: '_cooperative_o_vs32',
KernelScheduleType.Mxf4TmaWarpSpecializedPingpongSm120: '_pingpong_o_vs32',
KernelScheduleType.F8f6f4SparseTmaWarpSpecializedCooperativeSm120: '_q'
KernelScheduleType.F8f6f4SparseTmaWarpSpecializedCooperativeSm120: '_q',
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperativeSm120: '_cooperative_q',
KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120: '_pingpong_q'
}
class EpilogueScheduleType(enum.Enum):
@@ -756,6 +774,20 @@ EpilogueFunctor3xTag = {
EpilogueFunctor3x.LinearCombinationBlockScaleFactor: 'cutlass::epilogue::fusion::LinCombBlockScaleFactor',
}
# TMA epilogues have certain alignment requirements as calculated in get_tma_alignment(data_type)
def is_tma_epilogue(epilogue_schedule_type):
return epilogue_schedule_type in [
EpilogueScheduleType.ScheduleAuto,
EpilogueScheduleType.TmaWarpSpecialized,
EpilogueScheduleType.TmaWarpSpecializedCooperative,
EpilogueScheduleType.TmaWarpSpecialized1Sm,
EpilogueScheduleType.TmaWarpSpecialized2Sm,
EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm,
EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm,
EpilogueScheduleType.PtrArrayTmaWarpSpecializedCooperative,
EpilogueScheduleType.PtrArrayTmaWarpSpecializedPingpong,
]
def to_grouped_schedule(schedule, grouped):
if not grouped:
return schedule
@@ -771,17 +803,18 @@ def to_grouped_schedule(schedule, grouped):
EpilogueScheduleType.TmaWarpSpecializedCooperative : EpilogueScheduleType.PtrArrayTmaWarpSpecializedCooperative,
EpilogueScheduleType.NoSmemWarpSpecialized : EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized,
# SM100
KernelScheduleType.TmaWarpSpecialized1SmSm100: KernelScheduleType.PtrArrayTmaWarpSpecialized1SmSm100,
KernelScheduleType.TmaWarpSpecialized2SmSm100: KernelScheduleType.PtrArrayTmaWarpSpecialized2SmSm100,
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100 : KernelScheduleType.PtrArrayNvf4TmaWarpSpecialized1SmSm100,
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100 : KernelScheduleType.PtrArrayNvf4TmaWarpSpecialized2SmSm100,
KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100 : KernelScheduleType.PtrArrayMxf4TmaWarpSpecialized1SmSm100,
KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100 : KernelScheduleType.PtrArrayMxf4TmaWarpSpecialized2SmSm100,
KernelScheduleType.Mxf8f6f4TmaWarpSpecialized1SmSm100 : KernelScheduleType.PtrArrayMxf8f6f4TmaWarpSpecialized1SmSm100,
KernelScheduleType.Mxf8f6f4TmaWarpSpecialized2SmSm100 : KernelScheduleType.PtrArrayMxf8f6f4TmaWarpSpecialized2SmSm100,
EpilogueScheduleType.TmaWarpSpecialized1Sm: EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm,
EpilogueScheduleType.TmaWarpSpecialized2Sm: EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm,
KernelScheduleType.BlockwiseTmaWarpSpecialized1SmSm100 : KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized1SmSm100,
KernelScheduleType.BlockwiseTmaWarpSpecialized2SmSm100 : KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized2SmSm100,
EpilogueScheduleType.TmaWarpSpecialized1Sm: EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm,
EpilogueScheduleType.TmaWarpSpecialized2Sm: EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm,
}
return group_schedule_map[schedule]

29
python/cutlass_library/manifest.py
View File

@@ -507,7 +507,8 @@ class Manifest:
self.selected_kernels = []
self.ignore_kernel_names = []
self.exclude_kernel_names = []
self.compute_capabilities = [50,]
self.compute_capabilities_baseline = [50,]
self.compute_capabilities_feature_set = ['50',]
self.curr_build_dir = '.'
self.filter_by_cc = True
@@ -518,21 +519,9 @@ class Manifest:
# A common user error is to use commas instead of semicolons.
if ',' in args.architectures:
raise RuntimeError("The list of architectures (CMake option CUTLASS_NVCC_ARCHS) must be semicolon-delimited.\nDon't use commas to separate the architectures; use semicolons.\nYou specified the list as: " + args.architectures)
architectures = args.architectures.split(';') if len(args.architectures) else ['50',]
arch_conditional_cc = [
'90a',
'100a',
'100f',
'101a',
'101f',
'120a',
'120f'
]
architectures = [x if x not in arch_conditional_cc else x.split('a')[0] for x in architectures]
architectures = [x if x not in arch_conditional_cc else x.split('f')[0] for x in architectures]
self.compute_capabilities = [int(x) for x in architectures]
self.compute_capabilities_feature_set = args.architectures.split(';') if len(args.architectures) else ['50',]
self.compute_capabilities_baseline = sorted(set(int(arch.split('a')[0].split('f')[0]) for arch in self.compute_capabilities_feature_set))
if args.filter_by_cc in ['false', 'False', '0']:
self.filter_by_cc = False
@@ -597,7 +586,7 @@ class Manifest:
return default_level
def get_kernel_filters (self, kernelListFile):
def get_kernel_filters(self, kernelListFile):
if os.path.isfile(kernelListFile):
with open(kernelListFile, 'r') as fileReader:
lines = [line.rstrip() for line in fileReader if not line.startswith("#")]
@@ -635,7 +624,7 @@ class Manifest:
# filter based on compute capability
enabled = not (self.filter_by_cc)
for cc in self.compute_capabilities:
for cc in self.compute_capabilities_baseline:
if cc >= operation.tile_description.minimum_compute_capability and \
cc <= operation.tile_description.maximum_compute_capability and \
@@ -789,14 +778,14 @@ class Manifest:
return name.endswith(".cpp")
def get_src_archs_str_given_requested_cuda_archs(archs, source_file):
intersected_archs = archs & set(self.compute_capabilities)
intersected_archs = archs & set(self.compute_capabilities_baseline)
if intersected_archs == set():
raise RuntimeError(
"""
Empty archs set for file {} after taking
the intersection of {} (global requested archs) and
{} (per file requested archs)
""".format(source_file, set(self.compute_capabilities), archs))
""".format(source_file, set(self.compute_capabilities_baseline), archs))
else:
return " ".join(map(str, intersected_archs))