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

v4.3 update. (#2709)

Browse Source 
* v4.3 update.

* Update the cute_dsl_api changelog's doc link

* Update version to 4.3.0

* Update the example link

* Update doc to encourage user to install DSL from requirements.txt

---------

Co-authored-by: Larry Wu <larwu@nvidia.com>
This commit is contained in:
Junkai-Wu
2025-10-21 14:26:30 -04:00
committed by GitHub
co-authored by Larry Wu
parent e6e2cc29f5
commit b1d6e2c9b3
244 changed files with 59272 additions and 10455 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/cutlass_library/emit_kernel_listing.py
-1
View File
@@ -594,7 +594,6 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
# reduce L1 test runtime if reference kernel is not running on device.
if mode == "functional_L1" and profiler_flags_for_verification == "host" :
problem_waves = [0.5, 2.5]
if dynamic_cluster:
if mode == "functional_L0":
python/cutlass_library/gemm_operation.py
+7 -3
View File
@@ -76,6 +76,7 @@ class GemmOperation:
GemmKind.GroupedBlockScaledUniversal3x,
GemmKind.BlockwiseUniversal3x,
GemmKind.GroupedBlockwiseUniversal3x,
GemmKind.BlockScaledSparseUniversal3x,
}
self.is_3x = gemm_kind in kinds_3x
self.prefix = "3x" if self.is_3x else ""
@@ -174,6 +175,7 @@ class GemmOperation:
OpcodeClass.WmmaTensorOp,
OpcodeClass.SparseTensorOp,
OpcodeClass.BlockScaledTensorOp,
OpcodeClass.BlockScaledSparseTensorOp,
]
is_tensor_op = self.tile_description.math_instruction.opcode_class in tensor_ops
@@ -348,7 +350,7 @@ class GemmOperation:
opcode_class_main = self.tile_description.math_instruction.opcode_class
instruction_shape = self.tile_description.math_instruction.instruction_shape
tile_shape_m, tile_shape_n, tile_shape_k = self.tile_description.tile_shape
if opcode_class_main in [OpcodeClass.TensorOp, OpcodeClass.BlockScaledTensorOp, OpcodeClass.SparseTensorOp]:
if opcode_class_main in [OpcodeClass.TensorOp, OpcodeClass.BlockScaledTensorOp, OpcodeClass.SparseTensorOp, OpcodeClass.BlockScaledSparseTensorOp]:
tile_shape_m = instruction_shape[0]
tile_shape_n = instruction_shape[1]
return (tile_shape_m, tile_shape_n, tile_shape_k)
@@ -984,7 +986,7 @@ ${compile_guard_end}
element_b = DataTypeTag[operation.B.element] if not operation.is_complex() else f"cute::tuple<{str(DataTypeTag[operation.B.element])},{str(ComplexTransformTag3x[operation.B.complex_transform])}>"
epilogue_schedule_type = EpilogueScheduleTag[operation.epilogue_schedule]
if opcode_class_main == OpcodeClass.BlockScaledTensorOp:
if opcode_class_main == OpcodeClass.BlockScaledTensorOp or opcode_class_main == OpcodeClass.BlockScaledSparseTensorOp:
grouped = is_grouped(operation.gemm_kind)
if cta_n == 256 and operation.kernel_schedule == to_grouped_schedule(KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100, grouped):
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
@@ -1099,7 +1101,7 @@ using {operation_name_str}_LayoutNarrowReordered = decltype(cute::tile_to_shape(
sfn_vec_size = operation.ScaleFactorNVecSize
sfk_vec_size = operation.ScaleFactorKVecSize
blockwise_prepare_code = f"""
using {operation_name_str}_ScaleConfig = cutlass::detail::Sm{operation.arch}BlockwiseScaleConfig<{sfm_vec_size}, {sfn_vec_size}, {sfk_vec_size}>;
using {operation_name_str}_ScaleConfig = cutlass::detail::Sm{"90" if operation.arch == 90 else "1xx"}BlockwiseScaleConfig<{sfm_vec_size}, {sfn_vec_size}, {sfk_vec_size}>;
using {operation_name_str}_LayoutSFA = decltype({operation_name_str}_ScaleConfig::deduce_layoutSFA());
using {operation_name_str}_LayoutSFB = decltype({operation_name_str}_ScaleConfig::deduce_layoutSFB());
"""
@@ -1477,6 +1479,7 @@ class EmitGemmConfigurationLibrary:
GemmKind.GroupedBlockScaledUniversal3x: EmitGemmUniversal3xInstance,
GemmKind.BlockwiseUniversal3x: EmitGemmUniversal3xInstance,
GemmKind.GroupedBlockwiseUniversal3x: EmitGemmUniversal3xInstance,
GemmKind.BlockScaledSparseUniversal3x: EmitGemmUniversal3xInstance,
}
self.gemm_kind_wrappers = {
@@ -1493,6 +1496,7 @@ class EmitGemmConfigurationLibrary:
GemmKind.GroupedBlockScaledUniversal3x: 'GroupedBlockScaledGemmUniversal3xOperation',
GemmKind.BlockwiseUniversal3x: 'BlockwiseGemmUniversal3xOperation',
GemmKind.GroupedBlockwiseUniversal3x: 'GroupedBlockwiseGemmUniversal3xOperation',
GemmKind.BlockScaledSparseUniversal3x: 'BlockScaledSparseGemmUniversal3xOperation',
}
self.wmma_guard_start = "#if defined(CUTLASS_ARCH_WMMA_SM${sm_number}_ENABLED)"
python/cutlass_library/generator.py
+701 -29
View File
@@ -6729,6 +6729,8 @@ try:
generate_f8f6f4_math_instructions_sm100,
generate_mxf8f6f4_math_instructions_sm100,
generate_mxf4nvf4_math_instructions_sm100,
generate_sparse_mxf4nvf4_math_instructions_sm100,
generate_sparse_mxf8f6f4_math_instructions_sm100,
generate_fp8_math_instructions_sm100,
generate_cluster_shapes_sm100,
get_pruning_level_from_global_level
@@ -6741,6 +6743,8 @@ except ImportError:
generate_f8f6f4_math_instructions_sm100,
generate_mxf8f6f4_math_instructions_sm100,
generate_mxf4nvf4_math_instructions_sm100,
generate_sparse_mxf4nvf4_math_instructions_sm100,
generate_sparse_mxf8f6f4_math_instructions_sm100,
generate_fp8_math_instructions_sm100,
generate_cluster_shapes_sm100,
get_pruning_level_from_global_level
@@ -6804,7 +6808,8 @@ def GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version):
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
math_instructions_1sm, math_instructions_2sm = generate_tf32_math_instructions_sm100(instantiation_level)
@@ -6879,7 +6884,9 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
math_instructions_1sm, math_instructions_2sm = generate_16b_math_instructions_sm100(instantiation_level)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
grouped = is_grouped(gemm_kind)
cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level)
@@ -7057,7 +7064,8 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
grouped = is_grouped(gemm_kind)
@@ -7505,7 +7513,8 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version, gemm_ki
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
@@ -7687,7 +7696,8 @@ def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cud
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
@@ -7918,7 +7928,8 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
@@ -8046,8 +8057,12 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
nvfp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100, grouped)
fp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100, grouped)
nvfp4_schedules = [[nvfp4_kernel_schedule, epi_schedule], [nvfp4_kernel_schedule, epi_nosmem_schedule]]
fp4_schedules = [[fp4_kernel_schedule, epi_schedule], [fp4_kernel_schedule, epi_nosmem_schedule]]
nvfp4_schedules = [[nvfp4_kernel_schedule, epi_schedule]]
fp4_schedules = [[fp4_kernel_schedule, epi_schedule]]
if (data_type["sfd_type"]["type"] == DataType.void):
nvfp4_schedules.append([nvfp4_kernel_schedule, epi_nosmem_schedule])
fp4_schedules.append([fp4_kernel_schedule, epi_nosmem_schedule])
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
)
@@ -8170,14 +8185,407 @@ def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
nvfp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100, grouped)
fp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100, grouped)
nvfp4_schedules = [[nvfp4_kernel_schedule, epi_schedule], [nvfp4_kernel_schedule, epi_nosmem_schedule]]
fp4_schedules = [[fp4_kernel_schedule, epi_schedule], [fp4_kernel_schedule, epi_nosmem_schedule]]
nvfp4_schedules = [[nvfp4_kernel_schedule, epi_schedule]]
fp4_schedules = [[fp4_kernel_schedule, epi_schedule]]
if (data_type["sfd_type"]["type"] == DataType.void):
nvfp4_schedules.append([nvfp4_kernel_schedule, epi_nosmem_schedule])
fp4_schedules.append([fp4_kernel_schedule, epi_nosmem_schedule])
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind)
if isFp4:
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind)
def GenerateSM100_SparseTensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledSparseUniversal3x):
# SM100 MMA with F4 inputs + block scale + sparse
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999)
grouped = is_grouped(gemm_kind)
if grouped:
return # not support for grouped sparse kernels
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.RowMajor, 64], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]],
[[LayoutType.RowMajor, 64], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]],
]
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
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]
math_instructions_1sm, math_instructions_2sm = generate_sparse_mxf4nvf4_math_instructions_sm100(instantiation_level)
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
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
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" : DataType.f32,
"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.bf16,
"d_type" : DataType.bf16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : DataType.f32,
"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.bf16,
"d_type" : DataType.e2m1,
"acc_type" : math_inst.element_accumulator,
"epi_type" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.e4m3, "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
nvfp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1SmNvf4
fp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1SmMxf4
nvfp4_kernel_schedule = KernelScheduleType.SparseNvf4TmaWarpSpecialized1SmSm100
fp4_kernel_schedule = KernelScheduleType.SparseMxf4TmaWarpSpecialized1SmSm100
nvfp4_schedules = [[nvfp4_kernel_schedule, nvfp4_epi_schedule]]
fp4_schedules = [[fp4_kernel_schedule, fp4_epi_schedule]]
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
)
if isFp4:
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
)
# 2xSM MMA kernels
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
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
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" : DataType.f32,
"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.bf16,
"d_type" : DataType.bf16,
"acc_type" : math_inst.element_accumulator,
"epi_type" : DataType.f32,
"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.bf16,
"d_type" : DataType.e2m1,
"acc_type" : math_inst.element_accumulator,
"epi_type" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.e4m3, "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
nvfp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2SmNvf4
fp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2SmMxf4
nvfp4_kernel_schedule = KernelScheduleType.SparseNvf4TmaWarpSpecialized2SmSm100
fp4_kernel_schedule = KernelScheduleType.SparseMxf4TmaWarpSpecialized2SmSm100
nvfp4_schedules = [[nvfp4_kernel_schedule, nvfp4_epi_schedule]]
fp4_schedules = [[fp4_kernel_schedule, fp4_epi_schedule]]
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
)
if isFp4:
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
)
def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledSparseUniversal3x):
# SM100 MMA with mixed F4/F6/F8 inputs + block scale + sparse
instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999)
grouped = is_grouped(gemm_kind)
if grouped:
return # not support for grouped sparse kernels
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 0]],
[[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.ColumnMajor, 0]],
[[LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 128], [LayoutType.RowMajor, 0]],
]
math_instructions_1sm, math_instructions_2sm = generate_sparse_mxf8f6f4_math_instructions_sm100(instantiation_level)
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)
# 1xSM MMA kernels
for math_inst in math_instructions_1sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
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
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] * 8 * 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" : DataType.f32,
"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.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor}
},
{
"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,
"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.bf16,
"d_type" : DataType.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor}
},
]
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
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1SmMxf8f6f4
kernel_schedule = KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized1SmSm100
schedules = [[kernel_schedule, epi_schedule]]
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
)
# 2xSM MMA kernels
for math_inst in math_instructions_2sm:
tile_descriptions = []
for cluster_shape in sm100_cluster_shape_1sm:
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
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] * 8 * 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" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
# void_c
{
"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" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor}
},
# none void_c
{
"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" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor}
},
{
"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,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
},
# void_c
{
"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" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor}
},
# none void_c
{
"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" : DataType.f32,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor}
},
]
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
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2SmMxf8f6f4
kernel_schedule = KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized2SmSm100
schedules = [[kernel_schedule, epi_schedule]]
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, schedules
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), 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
if not CudaToolkitVersionSatisfies(cuda_version, 13, 0):
@@ -8581,8 +8989,8 @@ def GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version):
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
epi_type = DataType.f32
math_instructions_1sm = [
@@ -8798,7 +9206,8 @@ def GenerateSM100_SparseTensorOp_32b_UMMA_gemm(manifest, cuda_version):
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
tile_schedulers = [
TileSchedulerType.Default, TileSchedulerType.StreamK
@@ -8926,7 +9335,8 @@ def GenerateSM100_SparseTensorOp_16b_UMMA_gemm(manifest, cuda_version):
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
tile_schedulers = [
TileSchedulerType.Default, TileSchedulerType.StreamK
@@ -9054,7 +9464,8 @@ def GenerateSM100_SparseTensorOp_int8_UMMA_gemm(manifest, cuda_version):
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
tile_schedulers = [
TileSchedulerType.Default, TileSchedulerType.StreamK
@@ -9181,7 +9592,8 @@ def GenerateSM100_SparseTensorOp_fp8_UMMA_gemm(manifest, cuda_version):
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
tile_schedulers = [
TileSchedulerType.Default, TileSchedulerType.StreamK
@@ -9322,7 +9734,8 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version):
thor_sm = ThorSMRenumbering(cuda_version)
min_cc = 100
max_cc = thor_sm
max_cc = 100
max_cc = max(max_cc, thor_sm)
tile_schedulers = [
TileSchedulerType.Default, TileSchedulerType.StreamK
@@ -9536,8 +9949,9 @@ def GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version,
thor_sm = ThorSMRenumbering(cuda_version)
minimum_compute_capability = 100
maximum_compute_capability = thor_sm
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
spatial_dims = [2, 3]
@@ -9584,7 +9998,7 @@ def GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version,
warp_count = [4, 1, 1]
tile_description = TileDescription(
tile_shape, stages, warp_count, math_inst,
minimum_compute_capability, maximum_compute_capability,
min_cc, max_cc,
cluster_shape)
tile_descriptions.append(tile_description)
@@ -9648,7 +10062,7 @@ def GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version,
warp_count = [4, 1, 1]
tile_description = TileDescription(
tile_shape, stages, warp_count, math_inst,
minimum_compute_capability, maximum_compute_capability,
min_cc, max_cc,
cluster_shape)
tile_descriptions.append(tile_description)
@@ -9691,8 +10105,10 @@ def GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version,
thor_sm = ThorSMRenumbering(cuda_version)
minimum_compute_capability = 100
maximum_compute_capability = thor_sm
min_cc = 100
max_cc = 100
max_cc = max(max_cc, thor_sm)
spatial_dims = [2, 3]
stages = 0 # zero means "deduce the number of stages automatically"
@@ -9732,7 +10148,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version,
warp_count = [4, 1, 1]
tile_description = TileDescription(
tile_shape, stages, warp_count, math_inst,
minimum_compute_capability, maximum_compute_capability,
min_cc, max_cc,
cluster_shape)
tile_descriptions.append(tile_description)
@@ -9797,7 +10213,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version,
warp_count = [4, 1, 1]
tile_description = TileDescription(
tile_shape, stages, warp_count, math_inst,
minimum_compute_capability, maximum_compute_capability,
min_cc, max_cc,
cluster_shape)
tile_descriptions.append(tile_description)
@@ -10138,6 +10554,252 @@ def GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio
gemm_kind = GemmKind.BlockScaledUniversal3x
)
def GenerateSM120_Sparse_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version):
# SM120 MMA with mixed F4/F6/F8 inputs + block scale
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
layouts = [
[[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 0]]
]
instruction_sizes = [
[16, 8, 64]
]
tile_sizes = [
[128, 64, 256],
[128, 128, 256]
]
cluster_shape = [1,1,1]
ab_types = [
# DataType.e2m1,
# DataType.e2m3,
# DataType.e3m2,
# DataType.e5m2,
DataType.e4m3,
]
acc_types = [ DataType.f32 ]
def tile_schedulers(sfdtype, kernel_schedule):
# Pingpong kernel schedule doesn't support stream-K.
# 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:
return [TileSchedulerType.Default]
else:
return [TileSchedulerType.Default, TileSchedulerType.StreamK]
min_cc = 120
max_cc = 121
epi_type = DataType.f32
math_instructions = []
kernel_schedules = [
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecializedAcc2x4Sm120
]
for instr_size, a_type, b_type, acc_type in product(instruction_sizes, ab_types, ab_types, acc_types):
math_instructions.append(
MathInstruction(
instr_size,
a_type, b_type, acc_type,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
for math_inst in math_instructions:
tile_descriptions = []
for tile_size in tile_sizes:
tile_descriptions.append(
TileDescription(tile_size, 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.f32,
"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.f16,
"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.e4m3,
"acc_type" : math_inst.element_accumulator,
"epi_type" : epi_type,
"sf_type" : math_inst.element_scale_factor,
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "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 data_type, kernel_schedule in product(data_types, kernel_schedules):
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, EpilogueScheduleType.SparseTmaWarpSpecializedCooperativeSm120]],
tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule),
gemm_kind = GemmKind.BlockScaledSparseUniversal3x
)
def GenerateSM120_Sparse_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version):
# SM120 MMA with with F4 + block scale
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
# layouts for ABC and their alignments.
layouts = [
[[LayoutType.RowMajor, 64], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]]
]
instruction_sizes = [
[16, 8, 64]
]
tile_sizes = [
[128, 64, 256],
[128, 128, 256]
]
cluster_shape = [1,1,1]
ab_types = [
DataType.e2m1
]
sf_types = [
DataType.ue4m3,
# DataType.ue8m0
]
acc_types = [ DataType.f32 ]
def is_nvf4(kernel_schedule):
if kernel_schedule == KernelScheduleType.SparseNvf4TmaWarpSpecializedSm120:
return True
else:
return False
def tile_schedulers(sfdtype, kernel_schedule):
# 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:
return [TileSchedulerType.Default]
else:
return [TileSchedulerType.Default, TileSchedulerType.StreamK]
min_cc = 120
max_cc = 121
epi_type = DataType.f32
math_instructions = []
kernel_schedules = [
KernelScheduleType.SparseNvf4TmaWarpSpecializedSm120,
KernelScheduleType.SparseMxf4TmaWarpSpecializedSm120,
]
for instr_size, a_type, b_type, acc_type, sf_type in product(instruction_sizes, ab_types, ab_types, acc_types, sf_types):
math_instructions.append(
MathInstruction(
instr_size,
a_type, b_type, acc_type,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
sf_type)
)
for math_inst in math_instructions:
for kernel_schedule in kernel_schedules:
tile_descriptions = []
for tile_size in tile_sizes:
# nvf4 kernel only supports ue4m3 SF
# mxf4 kernel only supports ue8m0 SF
if (math_inst.element_scale_factor == DataType.ue4m3 and is_nvf4(kernel_schedule)) or \
(math_inst.element_scale_factor == DataType.ue8m0 and not is_nvf4(kernel_schedule)):
tile_descriptions.append(
TileDescription(tile_size, 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.f32,
"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.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.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 data_type in data_types:
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[kernel_schedule, EpilogueScheduleType.SparseTmaWarpSpecializedCooperativeSm120]],
tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule),
gemm_kind = GemmKind.BlockScaledSparseUniversal3x
)
def GenerateSM120_Sparse_TensorOp_gemm(manifest, cuda_version):
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
return
@@ -10423,6 +11085,12 @@ def GenerateSM100(manifest, cuda_version):
GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version)
GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockScaledUniversal3x)
#
# Block Scaled Sparse Gemm
#
GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version)
GenerateSM100_SparseTensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version)
#
# Conv
#
@@ -10437,7 +11105,11 @@ def GenerateSM120(manifest, cuda_version):
#
GenerateSM120_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version)
GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version)
#
# Sparse Block Scaled Gemm
#
GenerateSM120_Sparse_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version)
GenerateSM120_Sparse_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version)
#
# Sparse Gemm
#
@@ -10460,8 +11132,8 @@ def GenerateSM90_Conv3x(manifest, cuda_version,
if not CudaToolkitVersionSatisfies(cuda_version, 12, 0):
return
minimum_compute_capability = 90
maximum_compute_capability = 90
min_cc = 90
max_cc = 90
spatial_dims = (2, 3)
@@ -10800,7 +11472,7 @@ def GenerateSM90_Conv3x(manifest, cuda_version,
math_inst = make_math_instruction(data_types, mma_shape)
tile_shape = (mma_shape[0], mma_shape[1], num_mma_per_tile * mma_shape[2])
tile_description = TileDescription(tile_shape, stages, warp_count, math_inst,
minimum_compute_capability, maximum_compute_capability, cluster_shape)
min_cc, max_cc, cluster_shape)
assert(isinstance(spatial_dim, int))
dims_and_alignments = (
(
python/cutlass_library/library.py
+71 -5
View File
@@ -322,7 +322,7 @@ def is_complex(data_type):
return False
def is_block_scaled(gemm_kind):
return gemm_kind in (GemmKind.BlockScaledUniversal3x, GemmKind.GroupedBlockScaledUniversal3x)
return gemm_kind in (GemmKind.BlockScaledUniversal3x, GemmKind.GroupedBlockScaledUniversal3x, GemmKind.BlockScaledSparseUniversal3x)
def is_blockwise(gemm_kind):
return gemm_kind in (GemmKind.BlockwiseUniversal3x, GemmKind.GroupedBlockwiseUniversal3x)
@@ -548,6 +548,12 @@ class KernelScheduleType(enum.Enum):
Nvf4TmaWarpSpecialized1SmSm100 = enum_auto()
Nvf4TmaWarpSpecialized2SmSm100 = enum_auto()
SparseMxf4TmaWarpSpecialized1SmSm100 = enum_auto()
SparseMxf4TmaWarpSpecialized2SmSm100 = enum_auto()
SparseNvf4TmaWarpSpecialized1SmSm100 = enum_auto()
SparseNvf4TmaWarpSpecialized2SmSm100 = enum_auto()
SparseMxf8f6f4TmaWarpSpecialized1SmSm100 = enum_auto()
SparseMxf8f6f4TmaWarpSpecialized2SmSm100 = enum_auto()
# FP4 Ultra
MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103 = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103 = enum_auto()
@@ -586,6 +592,10 @@ class KernelScheduleType(enum.Enum):
Mxf4TmaWarpSpecializedCooperativeSm120 = enum_auto()
Mxf4TmaWarpSpecializedPingpongSm120 = enum_auto()
SparseMxf8f6f4TmaWarpSpecializedSm120 = enum_auto()
SparseMxf8f6f4TmaWarpSpecializedAcc2x4Sm120 = enum_auto()
SparseNvf4TmaWarpSpecializedSm120 = enum_auto()
SparseMxf4TmaWarpSpecializedSm120 = enum_auto()
F8f6f4SparseTmaWarpSpecializedCooperativeSm120 = enum_auto()
BlockwiseTmaWarpSpecializedCooperativeSm120 = enum_auto()
@@ -637,6 +647,13 @@ KernelScheduleTag = {
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100',
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100',
KernelScheduleType.SparseMxf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelSparseTmaWarpSpecialized1SmMxf4Sm100',
KernelScheduleType.SparseMxf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelSparseTmaWarpSpecialized2SmMxf4Sm100',
KernelScheduleType.SparseNvf4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelSparseTmaWarpSpecialized1SmNvf4Sm100',
KernelScheduleType.SparseNvf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelSparseTmaWarpSpecialized2SmNvf4Sm100',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelSparseTmaWarpSpecialized1SmMxf8f6f4Sm100',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelSparseTmaWarpSpecialized2SmMxf8f6f4Sm100',
# FP4 Ultra
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103',
@@ -694,6 +711,10 @@ KernelScheduleTag = {
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperativeSm120: 'cutlass::gemm::KernelTmaWarpSpecializedBlockwiseCooperativeSm120',
KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120: 'cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecializedSm120: 'cutlass::gemm::KernelSparseTmaWarpSpecializedMxf8f6f4Sm120',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecializedAcc2x4Sm120: 'cutlass::gemm::KernelSparseTmaWarpSpecializedMxf8f6f4Acc2x4Sm120',
KernelScheduleType.SparseNvf4TmaWarpSpecializedSm120: 'cutlass::gemm::KernelSparseTmaWarpSpecializedNvf4Sm120',
KernelScheduleType.SparseMxf4TmaWarpSpecializedSm120: 'cutlass::gemm::KernelSparseTmaWarpSpecializedMxf4Sm120',
}
#
@@ -742,6 +763,14 @@ KernelScheduleSuffixes = {
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: '_o_vs16_1sm',
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100: '_o_vs16_2sm',
KernelScheduleType.SparseMxf4TmaWarpSpecialized1SmSm100: '_o_vs32_1sm',
KernelScheduleType.SparseMxf4TmaWarpSpecialized2SmSm100: '_o_vs32_2sm',
KernelScheduleType.SparseNvf4TmaWarpSpecialized1SmSm100: '_o_vs16_1sm',
KernelScheduleType.SparseNvf4TmaWarpSpecialized2SmSm100: '_o_vs16_2sm',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized1SmSm100: '_q_1sm',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized2SmSm100: '_q_2sm',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: '_o_vs16_ultra_1sm',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: '_o_vs16_ultra_2sm',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: '_o_vs32_ultra_1sm',
@@ -796,6 +825,11 @@ KernelScheduleSuffixes = {
KernelScheduleType.Mxf4TmaWarpSpecializedCooperativeSm120: '_cooperative_o_vs32',
KernelScheduleType.Mxf4TmaWarpSpecializedPingpongSm120: '_pingpong_o_vs32',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecializedSm120: '_q',
KernelScheduleType.SparseMxf8f6f4TmaWarpSpecializedAcc2x4Sm120: '_acc2x4_q',
KernelScheduleType.SparseNvf4TmaWarpSpecializedSm120: '_o_vs16',
KernelScheduleType.SparseMxf4TmaWarpSpecializedSm120: '_o_vs32',
KernelScheduleType.F8f6f4SparseTmaWarpSpecializedCooperativeSm120: '_q',
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperativeSm120: '_cooperative_q',
@@ -827,6 +861,13 @@ class EpilogueScheduleType(enum.Enum):
PtrArrayTmaWarpSpecialized2Sm = enum_auto()
PtrArrayTmaWarpSpecializedPingpong = enum_auto()
PtrArrayTmaWarpSpecializedCooperative = enum_auto()
TmaWarpSpecialized1SmNvf4 = enum_auto()
TmaWarpSpecialized2SmNvf4 = enum_auto()
TmaWarpSpecialized1SmMxf4 = enum_auto()
TmaWarpSpecialized2SmMxf4 = enum_auto()
TmaWarpSpecialized1SmMxf8f6f4 = enum_auto()
TmaWarpSpecialized2SmMxf8f6f4 = enum_auto()
SparseTmaWarpSpecializedCooperativeSm120 = enum_auto()
#
EpilogueScheduleTag = {
@@ -854,6 +895,13 @@ EpilogueScheduleTag = {
EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm: 'cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm',
EpilogueScheduleType.PtrArrayTmaWarpSpecializedCooperative: 'cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative',
EpilogueScheduleType.PtrArrayTmaWarpSpecializedPingpong: 'cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong',
EpilogueScheduleType.TmaWarpSpecialized1SmNvf4: 'cutlass::epilogue::TmaWarpSpecialized1SmNvf4',
EpilogueScheduleType.TmaWarpSpecialized2SmNvf4: 'cutlass::epilogue::TmaWarpSpecialized2SmNvf4',
EpilogueScheduleType.TmaWarpSpecialized1SmMxf4: 'cutlass::epilogue::TmaWarpSpecialized1SmMxf4',
EpilogueScheduleType.TmaWarpSpecialized2SmMxf4: 'cutlass::epilogue::TmaWarpSpecialized2SmMxf4',
EpilogueScheduleType.TmaWarpSpecialized1SmMxf8f6f4: 'cutlass::epilogue::TmaWarpSpecialized1SmMxf8f6f4',
EpilogueScheduleType.TmaWarpSpecialized2SmMxf8f6f4: 'cutlass::epilogue::TmaWarpSpecialized2SmMxf8f6f4',
EpilogueScheduleType.SparseTmaWarpSpecializedCooperativeSm120: 'cutlass::epilogue::SparseTmaWarpSpecializedCooperativeSm120',
}
#
@@ -882,6 +930,13 @@ EpilogueScheduleSuffixes = {
EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm: '_epi_tma',
EpilogueScheduleType.PtrArrayTmaWarpSpecializedCooperative: '_epi_tma',
EpilogueScheduleType.PtrArrayTmaWarpSpecializedPingpong: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized1SmNvf4: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized2SmNvf4: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized1SmMxf4: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized2SmMxf4: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized1SmMxf8f6f4: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized2SmMxf8f6f4: '_epi_tma',
EpilogueScheduleType.SparseTmaWarpSpecializedCooperativeSm120: '_epi_tma',
}
class EpilogueFunctor3x(enum.Enum):
@@ -906,6 +961,12 @@ def is_tma_epilogue(epilogue_schedule_type):
EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm,
EpilogueScheduleType.PtrArrayTmaWarpSpecializedCooperative,
EpilogueScheduleType.PtrArrayTmaWarpSpecializedPingpong,
EpilogueScheduleType.TmaWarpSpecialized1SmNvf4,
EpilogueScheduleType.TmaWarpSpecialized2SmNvf4,
EpilogueScheduleType.TmaWarpSpecialized1SmMxf4,
EpilogueScheduleType.TmaWarpSpecialized2SmMxf4,
EpilogueScheduleType.TmaWarpSpecialized1SmMxf8f6f4,
EpilogueScheduleType.TmaWarpSpecialized2SmMxf8f6f4,
]
def to_grouped_schedule(schedule, grouped):
@@ -1040,7 +1101,8 @@ class OpcodeClass(enum.Enum):
TensorOp = enum_auto()
WmmaTensorOp = enum_auto()
SparseTensorOp = enum_auto()
BlockScaledTensorOp = enum_auto()
BlockScaledTensorOp = enum_auto()
BlockScaledSparseTensorOp = enum_auto()
OpcodeClassNames = {
@@ -1048,7 +1110,8 @@ OpcodeClassNames = {
OpcodeClass.TensorOp: 'tensorop',
OpcodeClass.WmmaTensorOp: 'wmma_tensorop',
OpcodeClass.SparseTensorOp: 'sptensorop',
OpcodeClass.BlockScaledTensorOp: 'bstensorop'
OpcodeClass.BlockScaledTensorOp: 'bstensorop',
OpcodeClass.BlockScaledSparseTensorOp: 'bssptensorop'
}
OpcodeClassTag = {
@@ -1056,7 +1119,8 @@ OpcodeClassTag = {
OpcodeClass.TensorOp: 'cutlass::arch::OpClassTensorOp',
OpcodeClass.WmmaTensorOp: 'cutlass::arch::OpClassWmmaTensorOp',
OpcodeClass.SparseTensorOp: 'cutlass::arch::OpClassSparseTensorOp',
OpcodeClass.BlockScaledTensorOp: 'cutlass::arch::OpClassBlockScaledTensorOp'
OpcodeClass.BlockScaledTensorOp: 'cutlass::arch::OpClassBlockScaledTensorOp',
OpcodeClass.BlockScaledSparseTensorOp: 'cutlass::arch::OpClassBlockScaledSparseTensorOp'
}
###################################################################################################
@@ -1143,6 +1207,7 @@ class GemmKind(enum.Enum):
GroupedBlockScaledUniversal3x = enum_auto()
BlockwiseUniversal3x = enum_auto()
GroupedBlockwiseUniversal3x = enum_auto()
BlockScaledSparseUniversal3x = enum_auto()
#
GemmKindNames = {
@@ -1158,7 +1223,8 @@ GemmKindNames = {
GemmKind.GroupedUniversal3x: "gemm_grouped",
GemmKind.GroupedBlockScaledUniversal3x: "gemm_grouped",
GemmKind.BlockwiseUniversal3x: "gemm",
GemmKind.GroupedBlockwiseUniversal3x: "gemm_grouped"
GemmKind.GroupedBlockwiseUniversal3x: "gemm_grouped",
GemmKind.BlockScaledSparseUniversal3x: "spgemm"
}
#
python/cutlass_library/sm100_shapes.py
+30
View File
@@ -315,6 +315,12 @@ SM100_MMA_SHAPES_MXF8F6F4_DENSE_1SM = {
}
SM100_MMA_SHAPES_MXF8F6F4_SPARSE_1SM = {
(128, 128, 32): 0,
(128, 192, 32): 1,
(128, 256, 32): 0,
}
SM100_MMA_SHAPES_MXF8F6F4_DENSE_2SM = {
(256, 64, 32): 1,
(256, 128, 32): 0,
@@ -324,6 +330,15 @@ SM100_MMA_SHAPES_MXF8F6F4_DENSE_2SM = {
}
SM100_MMA_SHAPES_MXF8F6F4_SPARSE_2SM = {
(256, 128, 32): 0,
(256, 192, 32): 1,
(256, 256, 32): 0,
}
# MXF4NVF4
SM100_MMA_SHAPES_MXF4NVF4_DENSE_1SM = {
(128, 64, 64): 1,
@@ -332,6 +347,13 @@ SM100_MMA_SHAPES_MXF4NVF4_DENSE_1SM = {
(128, 256, 64): 0,
}
SM100_MMA_SHAPES_MXF4NVF4_SPARSE_1SM = {
(128, 128, 64): 0,
(128, 256, 64): 0,
}
SM100_MMA_SHAPES_MXF4NVF4_DENSE_2SM = {
# Multiples of 16 for N
(256, 64, 64): 1,
@@ -340,3 +362,11 @@ SM100_MMA_SHAPES_MXF4NVF4_DENSE_2SM = {
(256, 256, 64): 0,
}
SM100_MMA_SHAPES_MXF4NVF4_SPARSE_2SM = {
# Multiples of 16 for N
(256, 128, 64): 0,
(256, 256, 64): 0,
}
python/cutlass_library/sm100_utils.py
+225
View File
@@ -659,3 +659,228 @@ def generate_cluster_shapes_sm100(level: int, change_priority_func : Union[Calla
]
return shapes_1sm, shapes_2sm
def generate_sparse_mxf4nvf4_math_instructions_sm100(level: int, enable_runtime_dtype = False, enable_compile_time_dtype = True):
"""
Generate all BlockScaledSparseTensorOp math instructions for MXFP4 and MXFP4 MMA that are supported by SM100 at or above the given level.
Args:
level: The global level to generate math instructions for.
enable_runtime_dtype: Whether to generate runtime dtype math instructions.
enable_compile_time_dtype: Whether to generate compile time dtype math instructions.
Returns:
A tuple of two lists of MathInstruction objects.
The first list contains the math instructions for 1SM, and the second list contains the math instructions for 2SM.
"""
tcgen05_level = get_tcgen05_level_from_global_level(level)
math_instructions_1sm = []
math_instructions_2sm = []
shapes_1sm = [
shape for shape, min_level in SM100_MMA_SHAPES_MXF4NVF4_SPARSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_MXF4NVF4_SPARSE_2SM.items() if tcgen05_level >= min_level
]
for shape in shapes_1sm:
if enable_runtime_dtype:
runtime_types = [ DataType.f4 ]
for a_type, b_type in product(runtime_types, repeat=2):
# math_instructions_1sm.append(
# MathInstruction(
# shape,
# a_type, b_type, DataType.f32,
# OpcodeClass.BlockScaledSparseTensorOp,
# MathOperation.multiply_add,
# DataType.ue8m0)
# )
math_instructions_1sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue4m3)
)
if enable_compile_time_dtype:
compile_time_types = [ DataType.e2m1,
]
for a_type, b_type in product(compile_time_types, repeat=2):
# math_instructions_1sm.append(
# MathInstruction(
# shape,
# a_type, b_type, DataType.f32,
# OpcodeClass.BlockScaledSparseTensorOp,
# MathOperation.multiply_add,
# DataType.ue8m0)
# )
math_instructions_1sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue4m3)
)
for shape in shapes_2sm:
if enable_runtime_dtype:
runtime_types = [ DataType.f4 ]
for a_type, b_type in product(runtime_types, repeat=2):
# math_instructions_2sm.append(
# MathInstruction(
# shape,
# a_type, b_type, DataType.f32,
# OpcodeClass.BlockScaledSparseTensorOp,
# MathOperation.multiply_add,
# DataType.ue8m0)
# )
math_instructions_2sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue4m3)
)
if enable_compile_time_dtype:
compile_time_types = [ DataType.e2m1,
]
for a_type, b_type in product(compile_time_types, repeat=2):
# math_instructions_2sm.append(
# MathInstruction(
# shape,
# a_type, b_type, DataType.f32,
# OpcodeClass.BlockScaledSparseTensorOp,
# MathOperation.multiply_add,
# DataType.ue8m0)
# )
math_instructions_2sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue4m3)
)
return math_instructions_1sm, math_instructions_2sm
def generate_sparse_mxf8f6f4_math_instructions_sm100(level: int, enable_runtime_dtype = False, enable_compile_time_dtype = True):
"""
Generate all BlockScaledSparseTensorOp math instructions for MXFP8, MXFP6, and MXFP4 MMA that are supported by SM100 at or above the given level.
Args:
level: The global level to generate math instructions for.
enable_runtime_dtype: Whether to generate runtime dtype math instructions.
enable_compile_time_dtype: Whether to generate compile time dtype math instructions.
Returns:
A tuple of two lists of MathInstruction objects.
The first list contains the math instructions for 1SM, and the second list contains the math instructions for 2SM.
"""
tcgen05_level = get_tcgen05_level_from_global_level(level)
pruning_level = get_pruning_level_from_global_level(level)
math_instructions_1sm = []
math_instructions_2sm = []
shapes_1sm = [
shape for shape, min_level in SM100_MMA_SHAPES_MXF8F6F4_SPARSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_MXF8F6F4_SPARSE_2SM.items() if tcgen05_level >= min_level
]
for shape in shapes_1sm:
if enable_runtime_dtype:
runtime_types = [ DataType.f8, DataType.f6, DataType.f4 ]
for a_type, b_type in product(runtime_types, repeat=2):
if pruning_level < 2 and ((a_type == DataType.f8 or b_type == DataType.f8)):
continue
math_instructions_1sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
if enable_compile_time_dtype:
compile_time_types = [ DataType.e4m3,
# DataType.e5m2,
# DataType.e3m2,
# DataType.e2m3,
# DataType.e2m1
]
for a_type, b_type in product(compile_time_types, repeat=2):
math_instructions_1sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
for shape in shapes_2sm:
if enable_runtime_dtype:
runtime_types = [ DataType.f8, DataType.f6, DataType.f4 ]
for a_type, b_type in product(runtime_types, repeat=2):
if pruning_level < 2 and ((a_type == DataType.f8 or b_type == DataType.f8)):
continue
math_instructions_2sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
if enable_compile_time_dtype:
compile_time_types = [ DataType.e4m3,
# DataType.e5m2,
# DataType.e3m2,
# DataType.e2m3,
# DataType.e2m1
]
for a_type, b_type in product(compile_time_types, repeat=2):
math_instructions_2sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledSparseTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
return math_instructions_1sm, math_instructions_2sm