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

v4.4 update. (#2979)

Browse Source
This commit is contained in:
Junkai-Wu
2026-01-24 11:46:17 -05:00
committed by GitHub
parent 2fafefb7b9
commit 9fba3195f9
293 changed files with 46343 additions and 2995 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/cutlass_library/emit_kernel_listing.py
+59 -4
View File
@@ -48,6 +48,7 @@ import csv
import json
import math
import os
import re
try:
import builtins
@@ -413,13 +414,49 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
raise Exception(error_message)
elif mode == "functional_L1":
sm100_mma_data_type_general = [
'gemm_f16_f16_f16_f16_f16',
'gemm_f16_f16_f16_void_f16',
'gemm_f16_f16_f32_f32_f32',
'gemm_bf16_bf16_f32_bf16_bf16',
'gemm_bf16_bf16_f32_f32_f32',
'gemm_e2m1_e2m1_f32_f32_f32',
'gemm_e2m1_e3m2_f32_f32_f32',
'gemm_e2m1_e4m3_f32_f32_f32',
'gemm_e3m2_e2m1_f32_f32_f32',
'gemm_e3m2_e3m2_f32_f32_f32',
'gemm_e3m2_e4m3_f32_f32_f32',
'gemm_e4m3_e2m1_f32_f32_f32',
'gemm_e4m3_e3m2_f32_f32_f32',
'gemm_e4m3_e4m3_f32_bf16_bf16',
'gemm_e4m3_e5m2_f32_bf16_e4m3',
'gemm_e5m2_e4m3_f32_f16_e4m3',
'gemm_s8_s8_s32_s32_s32',
'gemm_s8_s8_s32_s8_s8',
'gemm_f4_f4_f32_f32_f32',
'gemm_f4_f6_f32_f32_f32',
'gemm_f4_f8_f32_f32_f32',
'gemm_f6_f4_f32_f32_f32',
'gemm_f6_f6_f32_f32_f32',
'gemm_f6_f8_f32_f32_f32',
'gemm_f8_f4_f32_f32_f32',
'gemm_f8_f6_f32_f32_f32',
'gemm_f8_f8_f32_bf16_bf16',
'gemm_f8_f8_f32_bf16_e4m3',
'gemm_f8_f8_f32_bf16_e5m2',
'gemm_f8_f8_f32_f16_e4m3',
'gemm_f8_f8_f32_f16_e5m2',
'gemm_f8_f8_f32_f16_f16',
'gemm_f8_f8_f32_f32_f32',
'tf32gemm_*',
]
sm100_mma_cluster_size = [
'0x0x1' # dynamic cluster
]
# Restrict to two layouts to reduce L1 build and test time.
sm100_mma_layouts = ['tnt', 'ntn']
sm100_mma_filter_regex_1sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*1sm.*"
sm100_mma_filter_regex_2sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [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.*"
block_scaled_data_type = [
'ue8m0xe2m1_ue8m0xe2m1_f32_f16_e5m2',
'ue8m0xe2m1_ue8m0xe2m3_f32_f16_e5m2',
@@ -446,7 +483,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
filter_regex_sm100_mma = f"({sm100_mma_filter_regex_1sm})|" \
f"({sm100_mma_filter_regex_2sm})|" \
f"({block_scaled_filter_regex_1sm})|" \
f"({block_scaled_filter_regex_2sm})" \
f"({block_scaled_filter_regex_2sm})|" \
f"({sm103_block_scaled_filter_regex_1sm})|" \
f"({sm103_block_scaled_filter_regex_2sm})"
# CTA tiles for sm120 MMA - only run one tile size to reduce build/test times
@@ -483,7 +520,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
filter_regex_sm120_mma = f"({filter_regex_sm120_mma_0})|({filter_regex_sm120_mma_1})|({filter_regex_sm120_mma_2})|({filter_regex_sm120_mma_3})"
problem_waves = [0.5, 1.25, 2.5]
problem_waves = [0.5, 2.5]
if arch in ["120a", "120f", "121a", "121f"]:
kernel_filter = f"({filter_regex_sm120_mma})"
@@ -538,6 +575,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
runtime_input_datatypes = [None]
if dynamic_datatype:
# Standard runtime datatype kernels encoded as f4_f4 / f6_f6 / f8_f8, etc.
if "f4_f4" in kernel_name:
runtime_input_datatypes = [['e2m1','e2m1']]
elif "f4_f6" in kernel_name:
@@ -588,6 +626,23 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
elif "ue8m0xf8_ue8m0xf8" in kernel_name:
runtime_input_datatypes = [['e4m3','e4m3']]
# Blockwise runtime-datatype kernels encode the fp8 selector together with the
# accumulator precision and block tile, e.g.:
# gemm_64x128f32xf8_32x128f32xf8_...
# which does not contain an "f8_f8" substring. As a fallback, detect this
# encoding and map it to the same runtime input datatypes as the symmetric
# f4_f4 / f6_f6 / f8_f8 cases above.
if runtime_input_datatypes == [None]:
m = re.search(r"f32x(f[468])", kernel_name)
if m:
fp_token = m.group(1)
if fp_token == "f4":
runtime_input_datatypes = [['e2m1', 'e2m1']]
elif fp_token == "f6":
runtime_input_datatypes = [['e3m2', 'e3m2']]
elif fp_token == "f8":
runtime_input_datatypes = [['e4m3', 'e4m3']]
if "bstensorop" in kernel_name or is_blockwise(manifest.operations_by_name[kernel_name].gemm_kind):
profiler_flags_for_verification = "host"
python/cutlass_library/generator.py
+230 -3
View File
@@ -7487,7 +7487,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kin
tile_schedulers = [
TileSchedulerType.Default
]
# Some SM100 NoSmem epilogue instantiations rely on CUTE's shape_div, which enforces a compile-time
# divisibility condition between CTA N and the epilogue N tile. Keep this conservative and scoped:
# only apply the divisibility filter for selected common (c_type, d_type) pairs.
@@ -7498,7 +7497,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kin
(DataType.void, DataType.f16): 64,
(DataType.void, DataType.bf16): 64,
}
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
@@ -7618,7 +7616,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kin
kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100
epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm
# SM100 NoSmem epilogue uses EpilogueTileAuto with N-tile = min(64, cta_n).
# CUTE's shape_div then requires a compile-time divisibility condition between cta_n and 64.
# Only instantiate kernels where cta_n <= 64 or cta_n is an exact multiple of 64 to avoid
@@ -10625,6 +10622,232 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
[[KernelScheduleType.SparseTmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]],
tile_schedulers=tile_schedulers)
# SM100 Interleaved Complex Tf32 Kernels
def GenerateSM100_TensorOp_32b_UMMA_gemm_complex(manifest, cuda_version):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 0):
return
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]],
[[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2]],
[[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]],
[[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2]]
]
complex_transforms = [
(ComplexTransform.none, ComplexTransform.none),
(ComplexTransform.conj, ComplexTransform.none),
(ComplexTransform.none, ComplexTransform.conj),
(ComplexTransform.conj, ComplexTransform.conj)
]
data_types = [
{
"a_type" : DataType.cf32,
"b_type" : DataType.cf32,
"c_type" : DataType.cf32,
"d_type" : DataType.cf32,
"acc_type" : DataType.cf32,
"epi_type" : DataType.cf32,
},
{
"a_type" : DataType.cf32,
"b_type" : DataType.cf32,
"c_type" : DataType.void,
"d_type" : DataType.cf32,
"acc_type" : DataType.cf32,
"epi_type" : DataType.cf32,
}
]
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
math_instructions_1sm = [
# tf32 -> f32
MathInstruction(
[128, 64, 4],
DataType.tf32, DataType.tf32, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add_complex)
]
cluster_shapes_1sm = [[1,2,1], [1,1,1], [1,4,1], [4,4,1]
, DynamicClusterShape
]
if thor_sm in manifest.compute_capabilities_baseline :
cluster_shapes_1sm = [[1,2,1], [1,1,1], [1,4,1]
, DynamicClusterShape
]
tile_schedulers = [
TileSchedulerType.Default, TileSchedulerType.StreamK
]
# 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))
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
[[KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized1SmSm100, EpilogueScheduleType.NoSmemWarpSpecialized1Sm]],
complex_transforms,
tile_schedulers=tile_schedulers)
# 2xSM MMA kernels
math_instructions_2sm = [
# tf32 -> f32
MathInstruction(
[256, 64, 4],
DataType.tf32, DataType.tf32, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add_complex)
]
cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1], [4,4,1]
, DynamicClusterShape
]
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
]
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in cluster_shapes_2sm:
multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2])
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_2sm[0],
math_inst.instruction_shape[1] * multiplier_2sm[1],
math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
[[KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized2SmSm100, EpilogueScheduleType.NoSmemWarpSpecialized2Sm]],
complex_transforms,
tile_schedulers=tile_schedulers)
def GenerateSM100_TensorOp_FastF32_UMMA_gemm_complex_stream_k(manifest, cuda_version):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 0):
return
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]],
[[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]],
[[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]],
]
data_types = [
{
"a_type" : DataType.cf32,
"b_type" : DataType.cf32,
"c_type" : DataType.cf32,
"d_type" : DataType.cf32,
"acc_type" : DataType.cf32,
"epi_type" : DataType.cf32,
}
]
# Unsupported yet
complex_transforms = None
# [
# (ComplexTransform.none, ComplexTransform.none),
# (ComplexTransform.conj, ComplexTransform.none),
# (ComplexTransform.none, ComplexTransform.conj),
# (ComplexTransform.conj, ComplexTransform.conj)
# ]
min_cc = 100
max_cc = 100
math_instructions_1sm = [
MathInstruction(
[128, 64, 8],
DataType.cbf16, DataType.cbf16, DataType.cf32,
OpcodeClass.TensorOp,
MathOperation.multiply_add),
]
cluster_shapes_1sm = [
[1,1,1], [4,4,1]
, DynamicClusterShape
]
tile_schedulers = [
TileSchedulerType.Default, TileSchedulerType.StreamK,
]
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in cluster_shapes_1sm:
multiplier = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier[0],
math_inst.instruction_shape[1] * multiplier[1],
math_inst.instruction_shape[2] * 2],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
[[KernelScheduleType.TmaWarpSpecialized1SmFastFP32Sm100, EpilogueScheduleType.FastF32NoSmemWarpSpecialized1Sm]],
complex_transforms,
tile_schedulers=tile_schedulers)
# 2xSM MMA kernels
math_instructions_2sm = [
MathInstruction(
[256, 64, 8],
DataType.cbf16, DataType.cbf16, DataType.cf32,
OpcodeClass.TensorOp,
MathOperation.multiply_add),
]
cluster_shapes_2sm = [
[2,1,1], [4,4,1]
, DynamicClusterShape
]
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in cluster_shapes_2sm:
multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2])
tile_descriptions.append(
TileDescription([
math_inst.instruction_shape[0] * multiplier_2sm[0],
math_inst.instruction_shape[1] * multiplier_2sm[1],
math_inst.instruction_shape[2] * 2],
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
[[KernelScheduleType.TmaWarpSpecialized2SmFastFP32Sm100, EpilogueScheduleType.FastF32NoSmemWarpSpecialized2Sm]],
complex_transforms,
tile_schedulers=tile_schedulers)
# Conv Utility functions
def make_dims_and_alignments_triple(dim: int, bit_per_element_A: int, bit_per_element_B: int, bit_per_element_C: int):
bit_alignment_required_by_tma = 128
@@ -11780,6 +12003,10 @@ def GenerateSM100(manifest, cuda_version):
GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version)
GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockwiseUniversal3x)
GenerateSM100_TensorOp_32b_UMMA_gemm_complex(manifest, cuda_version)
# CGemm with 9xBF16
GenerateSM100_TensorOp_FastF32_UMMA_gemm_complex_stream_k(manifest, cuda_version)
#
# Sparse Gemm
#
python/cutlass_library/library.py
+14 -2
View File
@@ -562,6 +562,12 @@ class KernelScheduleType(enum.Enum):
SparseNvf4TmaWarpSpecialized2SmSm100 = enum_auto()
SparseMxf8f6f4TmaWarpSpecialized1SmSm100 = enum_auto()
SparseMxf8f6f4TmaWarpSpecialized2SmSm100 = enum_auto()
InterleavedComplexTF32TmaWarpSpecialized1SmSm100 = enum_auto()
InterleavedComplexTF32TmaWarpSpecialized2SmSm100 = enum_auto()
TmaWarpSpecialized1SmFastFP32Sm100 = enum_auto()
TmaWarpSpecialized2SmFastFP32Sm100 = enum_auto()
# FP4 Ultra
MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103 = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103 = enum_auto()
@@ -679,7 +685,10 @@ KernelScheduleTag = {
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmInterleavedComplexTF32Sm100',
KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmInterleavedComplexTF32Sm100',
KernelScheduleType.TmaWarpSpecialized1SmFastFP32Sm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmFastFP32Sm100',
KernelScheduleType.TmaWarpSpecialized2SmFastFP32Sm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmFastFP32Sm100',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperative: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperativeFP8FastAccum: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum',
KernelScheduleType.PtrArrayTmaWarpSpecializedPingpong: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong',
@@ -799,7 +808,10 @@ KernelScheduleSuffixes = {
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: '_o_vs16_ultra_2sm_tmapf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_1sm_tmapf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_2sm_tmapf',
KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized1SmSm100: '_1sm',
KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized2SmSm100: '_2sm',
KernelScheduleType.TmaWarpSpecialized1SmFastFP32Sm100: '_FastF32_1sm',
KernelScheduleType.TmaWarpSpecialized2SmFastFP32Sm100: '_FastF32_2sm',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperative: '_warpspecialized_cooperative',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperativeFP8FastAccum: '_warpspecialized_cooperative_fp8_fastaccum',
KernelScheduleType.PtrArrayTmaWarpSpecializedPingpong: '_warpspecialized_pingpong',