v4.3 tag release update. (#2789)
This commit is contained in:
Junkai-Wu
@@ -5841,7 +5841,24 @@ def GenerateSM90_TensorOp_fp8_WGMMA_gemm_with_blockwise(manifest, cuda_version,
|
||||
level=instantiation_level)
|
||||
|
||||
tile_descriptions = list()
|
||||
|
||||
tile_descriptions.append(
|
||||
TileDescription(
|
||||
threadblock_shape=[
|
||||
256,
|
||||
128,
|
||||
128
|
||||
],
|
||||
stages=0,
|
||||
warp_count=[4, 1, 1],
|
||||
math_instruction=MathInstruction(
|
||||
[128, 128, 32],
|
||||
DataType.e5m2, DataType.e4m3, DataType.f32,
|
||||
OpcodeClass.TensorOp,
|
||||
MathOperation.multiply_add),
|
||||
min_compute=90,
|
||||
max_compute=90,
|
||||
cluster_shape=[1,2,1],
|
||||
explicit_vector_sizes=[1, 128, 128]))
|
||||
for desc in tile_descriptions_:
|
||||
desc.explicit_vector_sizes = [1, desc.tile_shape[1], desc.tile_shape[2]]
|
||||
tile_descriptions.append(copy.deepcopy(desc))
|
||||
@@ -7043,6 +7060,104 @@ def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed,
|
||||
[[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
def GenerateSM100_TensorOp_16b_UMMA_alignx_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x):
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
|
||||
return
|
||||
|
||||
instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999)
|
||||
|
||||
# layouts for ABC and their alignments. C alignment will be set later based on output type
|
||||
layouts = [
|
||||
[[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]],
|
||||
[[LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]],
|
||||
[[LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]],
|
||||
[[LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]],
|
||||
[[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]],
|
||||
[[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]],
|
||||
[[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]],
|
||||
[[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]],
|
||||
]
|
||||
|
||||
thor_sm = ThorSMRenumbering(cuda_version)
|
||||
|
||||
math_instructions_1sm, _ = generate_16b_math_instructions_sm100(instantiation_level)
|
||||
|
||||
min_cc = 100
|
||||
max_cc = 100
|
||||
max_cc = max(max_cc, thor_sm)
|
||||
|
||||
grouped = is_grouped(gemm_kind)
|
||||
if grouped:
|
||||
return
|
||||
cluster_shapes_1sm= [[1,1,1]]
|
||||
|
||||
tile_schedulers = [
|
||||
TileSchedulerType.Default
|
||||
]
|
||||
|
||||
# 1xSM MMA kernels
|
||||
for math_inst in math_instructions_1sm:
|
||||
tile_descriptions = []
|
||||
for cluster_shape in cluster_shapes_1sm:
|
||||
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
|
||||
tile_descriptions.append(
|
||||
TileDescription([
|
||||
math_inst.instruction_shape[0] * multiplier_1sm[0],
|
||||
math_inst.instruction_shape[1] * multiplier_1sm[1],
|
||||
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
|
||||
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
|
||||
|
||||
data_types = [
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : math_inst.element_accumulator,
|
||||
"d_type" : math_inst.element_accumulator,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : math_inst.element_accumulator,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : math_inst.element_accumulator,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : math_inst.element_accumulator,
|
||||
},
|
||||
]
|
||||
|
||||
kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100
|
||||
epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
|
||||
[[kernel_schedule, epi_schedule]],
|
||||
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
# for mixed precision kernels, also generate kernels that write output matrix in the A/B format
|
||||
# Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation)
|
||||
if math_inst.element_a != math_inst.element_accumulator:
|
||||
data_types_mixed = [
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : math_inst.element_a,
|
||||
"d_type" : math_inst.element_a,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : math_inst.element_accumulator,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : math_inst.element_a,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : math_inst.element_accumulator,
|
||||
},
|
||||
]
|
||||
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed,
|
||||
[[kernel_schedule, epi_schedule]],
|
||||
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x):
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
|
||||
return
|
||||
@@ -7198,16 +7313,19 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
|
||||
# Set alignment d based on Destination format.
|
||||
for layout in layouts:
|
||||
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
|
||||
|
||||
for data_type in data_types:
|
||||
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
|
||||
( data_type["d_type"] == DataType.e5m2 ):
|
||||
continue
|
||||
kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized1SmSm100, grouped)
|
||||
epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
|
||||
[[kernel_schedule, epi_schedule]],
|
||||
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
for tile_description in tile_descriptions:
|
||||
for layout in layouts:
|
||||
for data_type in data_types:
|
||||
if layout[1][0] == LayoutType.RowMajor and tile_description.math_instruction.instruction_shape[1] % 16 != 0:
|
||||
continue
|
||||
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
|
||||
( data_type["d_type"] == DataType.e5m2 ):
|
||||
continue
|
||||
kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized1SmSm100, grouped)
|
||||
epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)
|
||||
CreateGemmUniversal3xOperator(manifest, [layout], [tile_description], data_type,
|
||||
[[kernel_schedule, epi_schedule]],
|
||||
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
# 2xSM MMA kernels
|
||||
|
||||
@@ -7341,6 +7459,158 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmK
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
|
||||
[[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x):
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
|
||||
return
|
||||
|
||||
instantiation_level = manifest.get_instantiation_level(pruned_level=591 , default_level=591 , exhaustive_level=9999)
|
||||
|
||||
# layouts for ABC and their alignments.
|
||||
layouts = [
|
||||
[[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], # TN Layout
|
||||
[[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], # TN Layout
|
||||
]
|
||||
|
||||
thor_sm = ThorSMRenumbering(cuda_version)
|
||||
|
||||
min_cc = 100
|
||||
max_cc = 100
|
||||
max_cc = max(max_cc, thor_sm)
|
||||
|
||||
epi_type = DataType.f32
|
||||
grouped = is_grouped(gemm_kind)
|
||||
|
||||
math_instructions_1sm, _ = generate_fp8_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped)
|
||||
|
||||
cluster_shapes_1sm = [[1,1,1]]
|
||||
|
||||
tile_schedulers = [
|
||||
TileSchedulerType.Default
|
||||
]
|
||||
|
||||
# 1xSM MMA kernels
|
||||
for math_inst in math_instructions_1sm:
|
||||
tile_descriptions = []
|
||||
for cluster_shape in cluster_shapes_1sm:
|
||||
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
|
||||
tile_descriptions.append(
|
||||
TileDescription([
|
||||
math_inst.instruction_shape[0] * multiplier_1sm[0],
|
||||
math_inst.instruction_shape[1] * multiplier_1sm[1],
|
||||
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
|
||||
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
|
||||
|
||||
data_types = [
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.f16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e4m3,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.bf16,
|
||||
"d_type" : DataType.bf16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.bf16,
|
||||
"d_type" : DataType.e4m3,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.bf16,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f32,
|
||||
"d_type" : DataType.f32,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.f16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.bf16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.f32,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.e4m3,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
for data_type in data_types:
|
||||
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
|
||||
( data_type["d_type"] == DataType.e5m2 ):
|
||||
continue
|
||||
|
||||
kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100
|
||||
epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
|
||||
[[kernel_schedule, epi_schedule]],
|
||||
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.BlockwiseUniversal3x):
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
|
||||
return
|
||||
@@ -7467,7 +7737,7 @@ def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version,
|
||||
continue
|
||||
|
||||
is_runtime_datatype_a = is_runtime_datatype(data_type["a_type"])
|
||||
is_runtime_datatype_b = is_runtime_datatype(data_type["d_type"])
|
||||
is_runtime_datatype_b = is_runtime_datatype(data_type["b_type"])
|
||||
|
||||
# A/B datatypes should be both static or dynamic
|
||||
if (is_runtime_datatype_a != is_runtime_datatype_b):
|
||||
@@ -8585,6 +8855,418 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifes
|
||||
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, schedules
|
||||
, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind
|
||||
)
|
||||
def GenerateSM100_TensorOp_16b_UMMA_moe_gemm(manifest, cuda_version, gemm_kind=GemmKind.MoeGroupedUniversal3x):
|
||||
# SM100 MOE GEMM
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 13, 0):
|
||||
return
|
||||
|
||||
instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999)
|
||||
# layouts for ABC and their alignments. C alignment will be set later based on output type
|
||||
layouts = [
|
||||
[[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 0]],
|
||||
]
|
||||
|
||||
thor_sm = ThorSMRenumbering(cuda_version)
|
||||
|
||||
math_instructions_1sm, math_instructions_2sm = generate_16b_math_instructions_sm100(instantiation_level)
|
||||
|
||||
min_cc = 100
|
||||
max_cc = 100
|
||||
max_cc = max(max_cc, thor_sm)
|
||||
|
||||
cluster_shapes= [[1,1,1]]
|
||||
|
||||
tile_schedulers = [
|
||||
TileSchedulerType.Default
|
||||
]
|
||||
# 1xSM MMA kernels
|
||||
for math_inst in math_instructions_1sm:
|
||||
tile_descriptions = []
|
||||
for cluster_shape in cluster_shapes:
|
||||
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
|
||||
tile_descriptions.append(
|
||||
TileDescription([
|
||||
math_inst.instruction_shape[0] * multiplier_1sm[0],
|
||||
math_inst.instruction_shape[1] * multiplier_1sm[1],
|
||||
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
|
||||
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
|
||||
data_types = [
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : math_inst.element_accumulator,
|
||||
"d_type" : math_inst.element_b,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : math_inst.element_accumulator,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : math_inst.element_b,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : math_inst.element_accumulator,
|
||||
},
|
||||
]
|
||||
# Set alignment d based on Destination format.
|
||||
for layout in layouts:
|
||||
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
|
||||
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100
|
||||
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types,
|
||||
[[kernel_schedule, epi_schedule]],
|
||||
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
def GenerateSM100_TensorOp_fp8_UMMA_moe_gemm(manifest, cuda_version, gemm_kind=GemmKind.MoeGroupedUniversal3x):
|
||||
# SM100 MOE GEMM
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 13, 0):
|
||||
return
|
||||
|
||||
instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999)
|
||||
# layouts for ABC and their alignments. C alignment will be set later based on output type
|
||||
layouts = [
|
||||
[[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]],
|
||||
]
|
||||
|
||||
thor_sm = ThorSMRenumbering(cuda_version)
|
||||
|
||||
math_instructions_1sm, math_instructions_2sm = generate_fp8_math_instructions_sm100(instantiation_level)
|
||||
|
||||
min_cc = 100
|
||||
max_cc = 100
|
||||
max_cc = max(max_cc, thor_sm)
|
||||
# only support 1x1x1 cluster shape
|
||||
cluster_shapes= [[1,1,1]]
|
||||
epi_type = DataType.f32
|
||||
tile_schedulers = [
|
||||
TileSchedulerType.Default
|
||||
]
|
||||
# 1xSM MMA kernels
|
||||
for math_inst in math_instructions_1sm:
|
||||
tile_descriptions = []
|
||||
for cluster_shape in cluster_shapes:
|
||||
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
|
||||
tile_descriptions.append(
|
||||
TileDescription([
|
||||
math_inst.instruction_shape[0] * multiplier_1sm[0],
|
||||
math_inst.instruction_shape[1] * multiplier_1sm[1],
|
||||
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
|
||||
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
|
||||
data_types = [
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.f16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e4m3,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.bf16,
|
||||
"d_type" : DataType.bf16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.bf16,
|
||||
"d_type" : DataType.e4m3,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
},
|
||||
]
|
||||
# Set alignment d based on Destination format.
|
||||
for layout in layouts:
|
||||
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
|
||||
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmSm100
|
||||
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm
|
||||
for data_type in data_types:
|
||||
if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\
|
||||
( data_type["d_type"] == DataType.e5m2 ):
|
||||
continue
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
|
||||
[[kernel_schedule, epi_schedule]],
|
||||
tile_schedulers=tile_schedulers, gemm_kind=gemm_kind)
|
||||
|
||||
def GenerateSM100_TensorOp_mixed_8bits_UMMA_moe_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledMoeGroupedUniversal3x):
|
||||
# SM100 moe GEMM with mixed F4/F6/F8 inputs + block scale
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 13, 0):
|
||||
return
|
||||
|
||||
instantiation_level = manifest.get_instantiation_level(pruned_level=590, default_level=590, exhaustive_level=9999)
|
||||
|
||||
grouped = is_grouped(gemm_kind)
|
||||
|
||||
layouts = [
|
||||
[[LayoutType.RowMajor, 0], [LayoutType.ColumnMajor, 0], [LayoutType.ColumnMajor, 0]],
|
||||
]
|
||||
|
||||
math_instructions_1sm, math_instructions_2sm = generate_mxf8f6f4_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped)
|
||||
|
||||
|
||||
cluster_shapes_1sm = [[1,1,1]]
|
||||
|
||||
acc_types = [ DataType.f32 ]
|
||||
|
||||
def tile_schedulers(sfdtype):
|
||||
# Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void,
|
||||
# the epilogue is the traditional linear combination, for which we already have tests with stream-K.
|
||||
if sfdtype["type"] == DataType.void or grouped:
|
||||
return [TileSchedulerType.Default]
|
||||
else:
|
||||
return [TileSchedulerType.Default, TileSchedulerType.StreamK]
|
||||
|
||||
thor_sm = ThorSMRenumbering(cuda_version)
|
||||
|
||||
min_cc = 100
|
||||
max_cc = 100
|
||||
max_cc = max(max_cc, thor_sm)
|
||||
|
||||
epi_type = DataType.f32
|
||||
|
||||
is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8)
|
||||
|
||||
|
||||
# 1xSM MMA kernels
|
||||
for math_inst in math_instructions_1sm:
|
||||
assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp
|
||||
if DataTypeSize[math_inst.element_a] != DataTypeSize[math_inst.element_b] or DataTypeSize[math_inst.element_a] == 6:
|
||||
continue
|
||||
tile_descriptions = []
|
||||
for cluster_shape in cluster_shapes_1sm:
|
||||
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
|
||||
tile_descriptions.append(
|
||||
TileDescription([
|
||||
math_inst.instruction_shape[0] * multiplier_1sm[0],
|
||||
math_inst.instruction_shape[1] * multiplier_1sm[1],
|
||||
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
|
||||
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
|
||||
|
||||
data_types = [
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.bf16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e3m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
|
||||
}]
|
||||
|
||||
|
||||
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100
|
||||
epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm
|
||||
for data_type in data_types:
|
||||
# Set alignment d based on Destination format.
|
||||
for layout in layouts:
|
||||
layout[0][1] = 128 // DataTypeSize[data_type["a_type"]]
|
||||
layout[1][1] = 128 // DataTypeSize[data_type["b_type"]]
|
||||
layout[2][1] = 128 // DataTypeSize[data_type["d_type"]]
|
||||
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
|
||||
schedules=[[kernel_schedule, epi_schedule]], tile_schedulers=[TileSchedulerType.Default], gemm_kind=gemm_kind)
|
||||
|
||||
def GenerateSM100_TensorOp_fp4_UMMA_MoE_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledMoeGroupedUniversal3x):
|
||||
# SM100 MoE GEMM with F4 + block scale
|
||||
if not CudaToolkitVersionSatisfies(cuda_version, 12, 8):
|
||||
return
|
||||
|
||||
instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999)
|
||||
|
||||
|
||||
# layouts for ABC and their alignments.
|
||||
layouts = [
|
||||
[[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]],
|
||||
]
|
||||
|
||||
shapes_1sm = [
|
||||
(128, 64, 64), (128, 128, 64), (128, 192, 64), (128, 256, 64)
|
||||
]
|
||||
math_instructions_1sm = []
|
||||
for shape in shapes_1sm:
|
||||
math_instructions_1sm.append(
|
||||
MathInstruction(
|
||||
shape,
|
||||
DataType.e2m1, DataType.e2m1, DataType.f32,
|
||||
OpcodeClass.BlockScaledTensorOp,
|
||||
MathOperation.multiply_add,
|
||||
DataType.ue8m0)
|
||||
)
|
||||
cluster_shapes_1sm = [[1,1,1]]
|
||||
|
||||
acc_types = [ DataType.f32 ] # Accumulator is always 32 bits for block scaled MMA instructions
|
||||
|
||||
|
||||
thor_sm = ThorSMRenumbering(cuda_version)
|
||||
|
||||
min_cc = 100
|
||||
max_cc = 100
|
||||
max_cc = max(max_cc, thor_sm)
|
||||
|
||||
epi_type = DataType.f32
|
||||
|
||||
is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8)
|
||||
|
||||
|
||||
# 1xSM MMA kernels
|
||||
for math_inst in math_instructions_1sm:
|
||||
assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp
|
||||
|
||||
tile_descriptions = []
|
||||
for cluster_shape in cluster_shapes_1sm:
|
||||
multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape
|
||||
tile_descriptions.append(
|
||||
TileDescription([
|
||||
math_inst.instruction_shape[0] * multiplier_1sm[0],
|
||||
math_inst.instruction_shape[1] * multiplier_1sm[1],
|
||||
math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]],
|
||||
0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape))
|
||||
assert math_inst.instruction_shape[2] * 4 == 256
|
||||
|
||||
data_types = [
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.bf16,
|
||||
"d_type" : DataType.bf16,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.e2m1,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e5m2,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.void,
|
||||
"d_type" : DataType.e2m1,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e2m1,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor}
|
||||
},
|
||||
{
|
||||
"a_type" : math_inst.element_a,
|
||||
"b_type" : math_inst.element_b,
|
||||
"c_type" : DataType.f16,
|
||||
"d_type" : DataType.e2m1,
|
||||
"acc_type" : math_inst.element_accumulator,
|
||||
"epi_type" : epi_type,
|
||||
"sf_type" : math_inst.element_scale_factor,
|
||||
"sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor}
|
||||
}
|
||||
]
|
||||
|
||||
# Set alignment d based on Destination format.
|
||||
for layout in layouts:
|
||||
layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]]
|
||||
|
||||
for layout in layouts:
|
||||
for data_type in data_types:
|
||||
if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor):
|
||||
data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout.
|
||||
if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor):
|
||||
continue
|
||||
|
||||
# E2M1 x E2M1, vector size 32, E8
|
||||
# E2M1 x E2M1, vector size 16, UE4M3
|
||||
isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1
|
||||
|
||||
kernel_schedule = KernelScheduleType.MixedTmaCpAsyncWarpSpecialized1SmBlockScaledSm100
|
||||
epi_schedule = EpilogueScheduleType.ScheduleAuto
|
||||
|
||||
CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, [[kernel_schedule, epi_schedule]]
|
||||
, tile_schedulers=[TileSchedulerType.Default], gemm_kind=gemm_kind
|
||||
)
|
||||
|
||||
|
||||
def GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledUniversal3x):
|
||||
# SM100 MMA with F4 + block scale
|
||||
@@ -11047,17 +11729,22 @@ def GenerateSM100(manifest, cuda_version):
|
||||
# Dense Gemm
|
||||
#
|
||||
GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version)
|
||||
|
||||
GenerateSM100_TensorOp_16b_UMMA_alignx_gemm(manifest, cuda_version)
|
||||
GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version)
|
||||
|
||||
if not bool(set(manifest.compute_capabilities_feature_set).intersection(arch_family_cc)):
|
||||
GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version)
|
||||
|
||||
GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version)
|
||||
GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version)
|
||||
# grouped GEMM
|
||||
GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x)
|
||||
GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x)
|
||||
|
||||
# MOE grouped GEMM
|
||||
GenerateSM100_TensorOp_16b_UMMA_moe_gemm(manifest, cuda_version)
|
||||
GenerateSM100_TensorOp_fp8_UMMA_moe_gemm(manifest, cuda_version)
|
||||
GenerateSM100_TensorOp_mixed_8bits_UMMA_moe_gemm_with_block_scaled(manifest, cuda_version)
|
||||
GenerateSM100_TensorOp_fp4_UMMA_MoE_gemm_with_block_scaled(manifest, cuda_version)
|
||||
# StreamK is included in regular generation
|
||||
GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user