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

v4.2 tag release. (#2638)

Browse Source
This commit is contained in:
Junkai-Wu
2025-09-15 12:21:53 -04:00
committed by GitHub
parent 56f0718a97
commit 6a35b4d22f
161 changed files with 14056 additions and 3793 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/cutlass_library/emit_kernel_listing.py
+9 -4
View File
@@ -348,11 +348,15 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
'gemm.*ue8m0xf6_ue8m0xf6_f32_f16_ue8m0xe3m2',
]
block_scaled_tile_k = ['x128_', 'x256_']
sm103_block_scaled_data_type = [
'gemm.*ue8m0xf4_ue8m0xf4_f32_f16_e5m2',
'gemm.*ue8m0xf4_ue8m0xf4_f32_f16_ue8m0xe2m1',
]
sm103_block_scaled_tile_k = ['x768_']
block_scaled_cluster_size = [
'4x4x1', '2x1x1',
'0x0x1' # dynamic cluster
@@ -360,11 +364,12 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode
block_scaled_layouts = ['tnt']
# regex list must be in kernel procedural name order
block_scaled_filter_regex_1sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type, block_scaled_cluster_size, block_scaled_layouts]]) + ").*1sm.*"
block_scaled_filter_regex_2sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type, block_scaled_cluster_size, block_scaled_layouts]]) + ").*2sm.*"
block_scaled_filter_regex_1sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type, block_scaled_tile_k, block_scaled_cluster_size, block_scaled_layouts]]) + ").*1sm.*"
block_scaled_filter_regex_2sm = "cutlass3x_sm100_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [block_scaled_data_type, block_scaled_tile_k, block_scaled_cluster_size, block_scaled_layouts]]) + ").*2sm.*"
sm103_block_scaled_filter_regex_1sm = "cutlass3x_sm103_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [sm103_block_scaled_data_type, block_scaled_cluster_size, block_scaled_layouts]]) + ").*1sm.*"
sm103_block_scaled_filter_regex_2sm = "cutlass3x_sm103_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [sm103_block_scaled_data_type, block_scaled_cluster_size, block_scaled_layouts]]) + ").*2sm.*"
sm103_block_scaled_prefetch_policy = ['tmapf']
sm103_block_scaled_filter_regex_1sm = "cutlass3x_sm103_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [sm103_block_scaled_data_type, sm103_block_scaled_tile_k, block_scaled_cluster_size, block_scaled_layouts]]) + ").*1sm.*(" + "|".join(sm103_block_scaled_prefetch_policy) + ").*"
sm103_block_scaled_filter_regex_2sm = "cutlass3x_sm103_bstensorop.*(" + ").*(".join([ "|".join(x) for x in [sm103_block_scaled_data_type, sm103_block_scaled_tile_k, block_scaled_cluster_size, block_scaled_layouts]]) + ").*2sm.*(" + "|".join(sm103_block_scaled_prefetch_policy) + ").*"
if arch in ["100a", "100f"]:
kernel_filter = f"({sm100_mma_filter_regex_1sm})|" \
python/cutlass_library/gemm_operation.py
+23 -18
View File
@@ -985,33 +985,38 @@ ${compile_guard_end}
epilogue_schedule_type = EpilogueScheduleTag[operation.epilogue_schedule]
if opcode_class_main == OpcodeClass.BlockScaledTensorOp:
is_no_smem_epilogue = operation.epilogue_schedule in [EpilogueScheduleType.NoSmemWarpSpecialized1Sm, EpilogueScheduleType.NoSmemWarpSpecialized2Sm]
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>"
if not is_no_smem_epilogue:
if is_tma_epilogue(operation.epilogue_schedule):
epilogue_schedule_type = EpilogueScheduleTag[to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)]
if cta_n == 256 and operation.kernel_schedule == to_grouped_schedule(KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100, grouped):
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
if not is_no_smem_epilogue:
if is_tma_epilogue(operation.epilogue_schedule):
epilogue_schedule_type = EpilogueScheduleTag[to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized2Sm, grouped)]
if cta_n == 256 and operation.kernel_schedule == KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103:
# SM103 FP4 Ultra
is_sm103_fp4_ultra_1sm_kernel_schedule = operation.kernel_schedule in [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch, grouped)
]
is_sm103_fp4_ultra_2sm_kernel_schedule = operation.kernel_schedule in [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch, grouped),
to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch, grouped)
]
if cta_n == 256 and is_sm103_fp4_ultra_1sm_kernel_schedule:
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
if not is_no_smem_epilogue:
epilogue_schedule_type = EpilogueScheduleTag[EpilogueScheduleType.TmaWarpSpecialized1Sm]
if cta_n == 256 and operation.kernel_schedule == KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103:
if is_tma_epilogue(operation.epilogue_schedule):
epilogue_schedule_type = EpilogueScheduleTag[to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)]
if cta_n == 256 and is_sm103_fp4_ultra_2sm_kernel_schedule:
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
if not is_no_smem_epilogue:
epilogue_schedule_type = EpilogueScheduleTag[EpilogueScheduleType.TmaWarpSpecialized2Sm]
if cta_n == 256 and operation.kernel_schedule == KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103:
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
if not is_no_smem_epilogue:
epilogue_schedule_type = EpilogueScheduleTag[EpilogueScheduleType.TmaWarpSpecialized1Sm]
if cta_n == 256 and operation.kernel_schedule == KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103:
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
if not is_no_smem_epilogue:
epilogue_schedule_type = EpilogueScheduleTag[EpilogueScheduleType.TmaWarpSpecialized2Sm]
if is_tma_epilogue(operation.epilogue_schedule):
epilogue_schedule_type = EpilogueScheduleTag[to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized2Sm, grouped)]
element_a = f'cute::tuple<{str(element_a)},{str(DataTypeTag[operation.ScaleFactorA])}>'
element_b = f'cute::tuple<{str(element_b)},{str(DataTypeTag[operation.ScaleFactorB])}>'
python/cutlass_library/generator.py
+339 -1197
View File
File diff suppressed because it is too large Load Diff
python/cutlass_library/heuristics_provider.py
+1
View File
@@ -41,6 +41,7 @@ import logging
import ctypes
import functools
try:
import builtins
if hasattr(builtins, "CUTLASS_IGNORE_PACKAGE") and CUTLASS_IGNORE_PACKAGE == True:
python/cutlass_library/library.py
+116 -50
View File
@@ -508,6 +508,8 @@ class KernelScheduleType(enum.Enum):
BlockwiseTmaWarpSpecializedCooperative = enum_auto()
PtrArrayBlockwiseTmaWarpSpecializedCooperative = enum_auto()
BlockwiseTmaWarpSpecializedPingpong = enum_auto()
PtrArrayBlockwiseTmaWarpSpecializedPingpong = enum_auto()
TmaWarpSpecialized1SmSm100 = enum_auto()
TmaWarpSpecialized2SmSm100 = enum_auto()
@@ -547,20 +549,35 @@ class KernelScheduleType(enum.Enum):
Nvf4TmaWarpSpecialized2SmSm100 = enum_auto()
# FP4 Ultra
BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103 = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103 = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103 = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103 = enum_auto()
MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103 = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103 = enum_auto()
MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103 = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103 = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch = enum_auto()
BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch = enum_auto()
MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103 = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103 = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103 = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103 = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch = enum_auto()
PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch = enum_auto()
Mxf8f6f4TmaWarpSpecializedCooperativeSm120 = enum_auto()
Mxf8f6f4TmaWarpSpecializedPingpongSm120 = enum_auto()
@@ -589,7 +606,8 @@ KernelScheduleTag = {
KernelScheduleType.TmaWarpSpecializedPingpongFP8FastAccum: 'cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum',
KernelScheduleType.ImplicitTmaWarpSpecializedSm90: 'cutlass::conv::KernelImplicitTmaWarpSpecializedSm90',
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperative: 'cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum',
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperative: 'cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise',
KernelScheduleType.BlockwiseTmaWarpSpecializedPingpong: 'cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8Blockwise',
KernelScheduleType.TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmSm100',
KernelScheduleType.TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmSm100',
@@ -620,27 +638,28 @@ KernelScheduleTag = {
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100',
# FP4 Ultra
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperative: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperativeFP8FastAccum: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum',
KernelScheduleType.PtrArrayTmaWarpSpecializedPingpong: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong',
KernelScheduleType.PtrArrayTmaWarpSpecializedPingpongFP8FastAccum: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum',
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecializedCooperative: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum',
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecializedCooperative: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise',
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecializedPingpong: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise',
KernelScheduleType.PtrArrayTmaWarpSpecialized1SmBlockScaledSm100: "cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmBlockScaledSm100",
KernelScheduleType.PtrArrayTmaWarpSpecialized2SmBlockScaledSm100: "cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmBlockScaledSm100",
@@ -651,6 +670,19 @@ KernelScheduleTag = {
KernelScheduleType.PtrArrayMxf8f6f4TmaWarpSpecialized1SmSm100: "cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf8f6f4Sm100",
KernelScheduleType.PtrArrayMxf8f6f4TmaWarpSpecialized2SmSm100: "cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmMxf8f6f4Sm100",
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch',
KernelScheduleType.Mxf8f6f4TmaWarpSpecializedCooperativeSm120: 'cutlass::gemm::KernelTmaWarpSpecializedMxf8f6f4Sm120',
KernelScheduleType.Mxf8f6f4TmaWarpSpecializedPingpongSm120: 'cutlass::gemm::KernelTmaWarpSpecializedPingpongMxf8f6f4Sm120',
KernelScheduleType.Nvf4TmaWarpSpecializedCooperativeSm120: 'cutlass::gemm::KernelTmaWarpSpecializedNvf4Sm120',
@@ -681,7 +713,8 @@ KernelScheduleSuffixes = {
KernelScheduleType.ImplicitTmaWarpSpecializedSm90: '_warpspecialized',
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperative: '_warpspecialized_cooperative',
KernelScheduleType.BlockwiseTmaWarpSpecializedPingpong: '_warpspecialized_pingpong',
KernelScheduleType.TmaWarpSpecialized1SmSm100: '_1sm',
KernelScheduleType.TmaWarpSpecialized2SmSm100: '_2sm',
@@ -709,20 +742,20 @@ KernelScheduleSuffixes = {
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100: '_o_vs16_1sm',
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100: '_o_vs16_2sm',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: '_o_vs16_1sm',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: '_o_vs16_2sm',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: '_o_vs32_1sm',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103: '_o_vs32_2sm',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: '_o_vs16_ultra_1sm',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: '_o_vs16_ultra_2sm',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: '_o_vs32_ultra_1sm',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103: '_o_vs32_ultra_2sm',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch: '_o_vs16_1sm_nopf',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: '_o_vs16_2sm_nopf',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: '_o_vs32_1sm_nopf',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: '_o_vs32_2sm_nopf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch: '_o_vs16_ultra_1sm_nopf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: '_o_vs16_ultra_2sm_nopf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: '_o_vs32_ultra_1sm_nopf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: '_o_vs32_ultra_2sm_nopf',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch: '_o_vs16_1sm_tmapf',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: '_o_vs16_2sm_tmapf',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: '_o_vs32_1sm_tmapf',
KernelScheduleType.BlockScaledMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: '_o_vs32_2sm_tmapf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch: '_o_vs16_ultra_1sm_tmapf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: '_o_vs16_ultra_2sm_tmapf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_1sm_tmapf',
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_2sm_tmapf',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperative: '_warpspecialized_cooperative',
KernelScheduleType.PtrArrayTmaWarpSpecializedCooperativeFP8FastAccum: '_warpspecialized_cooperative_fp8_fastaccum',
@@ -730,6 +763,7 @@ KernelScheduleSuffixes = {
KernelScheduleType.PtrArrayTmaWarpSpecializedPingpongFP8FastAccum: '_warpspecialized_pingpong_fp8_fastaccum',
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecializedCooperative: '_warpspecialized_cooperative',
KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecializedPingpong: '_warpspecialized_pingpong',
KernelScheduleType.PtrArrayTmaWarpSpecialized1SmBlockScaledSm100: '_1sm',
KernelScheduleType.PtrArrayTmaWarpSpecialized2SmBlockScaledSm100: '_2sm',
@@ -740,6 +774,21 @@ KernelScheduleSuffixes = {
KernelScheduleType.PtrArrayMxf8f6f4TmaWarpSpecialized1SmSm100: '_o_vs32_1sm',
KernelScheduleType.PtrArrayMxf8f6f4TmaWarpSpecialized2SmSm100: '_o_vs32_2sm',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: '_o_vs16_ultra_1sm',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: '_o_vs16_ultra_2sm',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: '_o_vs32_ultra_1sm',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103: '_o_vs32_ultra_2sm',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch: '_o_vs16_ultra_1sm_nopf',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: '_o_vs16_ultra_2sm_nopf',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: '_o_vs32_ultra_1sm_nopf',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: '_o_vs32_ultra_2sm_nopf',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch: '_o_vs16_ultra_1sm_tmapf',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: '_o_vs16_ultra_2sm_tmapf',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_1sm_tmapf',
KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_2sm_tmapf',
KernelScheduleType.Mxf8f6f4TmaWarpSpecializedCooperativeSm120: '_cooperative_q',
KernelScheduleType.Mxf8f6f4TmaWarpSpecializedPingpongSm120: '_pingpong_q',
KernelScheduleType.Nvf4TmaWarpSpecializedCooperativeSm120: '_cooperative_o_vs16',
@@ -817,8 +866,8 @@ EpilogueScheduleSuffixes = {
EpilogueScheduleType.TmaWarpSpecializedCooperative: '_epi_tma',
EpilogueScheduleType.TmaWarpSpecialized1Sm: '',
EpilogueScheduleType.TmaWarpSpecialized2Sm: '_epi_tma',
EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm: '_tma_1sm',
EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm: '_tma_2sm',
EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm: '',
EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm: '_epi_tma',
EpilogueScheduleType.PtrArrayTmaWarpSpecializedCooperative: '_epi_tma',
EpilogueScheduleType.PtrArrayTmaWarpSpecializedPingpong: '_epi_tma',
}
@@ -855,6 +904,7 @@ def to_grouped_schedule(schedule, grouped):
# SM90
KernelScheduleType.TmaWarpSpecializedCooperative : KernelScheduleType.PtrArrayTmaWarpSpecializedCooperative,
KernelScheduleType.BlockwiseTmaWarpSpecializedCooperative : KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecializedCooperative,
KernelScheduleType.BlockwiseTmaWarpSpecializedPingpong : KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecializedPingpong,
KernelScheduleType.TmaWarpSpecializedPingpong : KernelScheduleType.PtrArrayTmaWarpSpecializedPingpong,
KernelScheduleType.TmaWarpSpecializedCooperativeFP8FastAccum : KernelScheduleType.PtrArrayTmaWarpSpecializedCooperativeFP8FastAccum,
KernelScheduleType.TmaWarpSpecializedPingpongFP8FastAccum : KernelScheduleType.PtrArrayTmaWarpSpecializedPingpongFP8FastAccum,
@@ -874,6 +924,21 @@ def to_grouped_schedule(schedule, grouped):
KernelScheduleType.BlockwiseTmaWarpSpecialized2SmSm100 : KernelScheduleType.PtrArrayBlockwiseTmaWarpSpecialized2SmSm100,
EpilogueScheduleType.TmaWarpSpecialized1Sm: EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm,
EpilogueScheduleType.TmaWarpSpecialized2Sm: EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm,
EpilogueScheduleType.NoSmemWarpSpecialized1Sm: EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized1Sm,
EpilogueScheduleType.NoSmemWarpSpecialized2Sm: EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized2Sm,
# SM103
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch,
KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: KernelScheduleType.PtrArrayMxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch,
}
return group_schedule_map[schedule]
@@ -1020,14 +1085,15 @@ ArchitectureNames = {
#
SharedMemPerCC = {
70: 96, # 96KB of SMEM
72: 96, # 96KB of SMEM
75: 64, # 64KB of SMEM
80: 163, # 163KB of SMEM - 1KB reserved for the driver
86: 99, # 99KB of SMEM - 1KB reserved for the driver
87: 163, # 163KB of SMEM - 1KB reserved for the driver
89: 99, # 99KB of SMEM - 1KB reserved for the driver
90: 227, # 227KB of SMEM - 1KB reserved for the driver
70: 96, # 96KB of SMEM
72: 96, # 96KB of SMEM
75: 64, # 64KB of SMEM
80: 163, # 163KB of SMEM - 1KB reserved for the driver
86: 99, # 99KB of SMEM - 1KB reserved for the driver
87: 163, # 163KB of SMEM - 1KB reserved for the driver
89: 99, # 99KB of SMEM - 1KB reserved for the driver
90: 227, # 227KB of SMEM - 1KB reserved for the driver
100: 227, # 227KB of SMEM - 1KB reserved for the driver
}
###################################################################################################
python/cutlass_library/manifest.py
+1 -1
View File
@@ -570,7 +570,7 @@ class Manifest:
self.kernel_filter_list.append(filter_re)
def get_sm90_instantiation_level(self, pruned_level=0, default_level=111, exhaustive_level=9992):
def get_instantiation_level(self, pruned_level=0, default_level=111, exhaustive_level=9992):
# Non-negative integer which determines how many kernels are instantiated.
# 0 = 0000 generates the fewest kernels, 9999 generates all possible combinations.
# increasing first digit reduces schedule / mixed type pruning,
python/cutlass_library/sm100_shapes.py
+342
View File
@@ -0,0 +1,342 @@
#################################################################################################
#
# Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#################################################################################################
"""
Valid tcgen05 shapes and cluster sizes for SM100, associated with levels.
These shape and level pairs are defined as dicts, where keys are shapes and values are their
associated levels. If the user input level for that category (tcgen05 shape, cluster
size) is smaller than a shape's associated level, it will be excluded, and otherwise, included.
Higher levels are therefore less likely emitted, but lower levels are more emitted more frequently.
Level 0 is always emitted.
"""
try:
from .library import DynamicClusterShape
except:
from library import DynamicClusterShape
SM100_CLUSTER_SHAPES_1SM = {
tuple(DynamicClusterShape) : 0,
# size 1 cluster
(1, 1, 1): 1,
# size 2 cluster
(1, 2, 1): 2,
(2, 1, 1): 5,
# size 4 clusters
(2, 2, 1): 6,
(1, 4, 1): 3,
(4, 1, 1): 6,
# size 8 clusters
(2, 4, 1): 7,
(4, 2, 1): 7,
(1, 8, 1): 8,
(8, 1, 1): 8,
# size 16 cluster
(4, 4, 1): 4,
}
SM100_CLUSTER_SHAPES_2SM = {
tuple(DynamicClusterShape) : 0,
# size 2 cluster
(2, 1, 1): 1,
# size 4 clusters
(2, 2, 1): 2,
(4, 1, 1): 2,
# size 8 clusters
(2, 4, 1): 3,
(4, 2, 1): 3,
(8, 1, 1): 6,
# size 16 cluster
(4, 4, 1): 4,
}
# MMA shapes
# 16b Dense
SM100_MMA_SHAPES_16b_DENSE_1SM = {
(64, 8, 16): 5,
(64, 16, 16): 2,
(64, 24, 16): 5,
(64, 32, 16): 2,
(64, 40, 16): 5,
(64, 48, 16): 5,
(64, 56, 16): 5,
(64, 64, 16): 2,
(64, 72, 16): 5,
(64, 80, 16): 5,
(64, 88, 16): 5,
(64, 96, 16): 5,
(64, 104, 16): 5,
(64, 112, 16): 5,
(64, 120, 16): 5,
(64, 128, 16): 0,
(64, 136, 16): 5,
(64, 144, 16): 5,
(64, 152, 16): 5,
(64, 160, 16): 5,
(64, 168, 16): 5,
(64, 176, 16): 5,
(64, 184, 16): 5,
(64, 192, 16): 3,
(64, 200, 16): 5,
(64, 208, 16): 5,
(64, 216, 16): 5,
(64, 224, 16): 5,
(64, 232, 16): 5,
(64, 240, 16): 5,
(64, 248, 16): 5,
(64, 256, 16): 3,
(128, 16, 16): 2,
(128, 32, 16): 2,
(128, 48, 16): 5,
(128, 64, 16): 2,
(128, 80, 16): 5,
(128, 96, 16): 5,
(128, 112, 16): 5,
(128, 128, 16): 0,
(128, 144, 16): 5,
(128, 160, 16): 5,
(128, 176, 16): 5,
(128, 192, 16): 3,
(128, 208, 16): 5,
(128, 224, 16): 5,
(128, 240, 16): 5,
(128, 256, 16): 0,
}
SM100_MMA_SHAPES_16b_DENSE_2SM = {
(128, 32, 16): 2,
(128, 64, 16): 2,
(128, 96, 16): 5,
(128, 128, 16): 0,
(128, 160, 16): 5,
(128, 192, 16): 5,
(128, 224, 16): 5,
(128, 256, 16): 0,
(256, 32, 16): 2,
(256, 64, 16): 2,
(256, 96, 16): 5,
(256, 128, 16): 0,
(256, 160, 16): 5,
(256, 192, 16): 3,
(256, 224, 16): 5,
(256, 256, 16): 0,
}
# TF32 Dense
SM100_MMA_SHAPES_TF32_DENSE_1SM = {
(64, 8, 8): 5,
(64, 16, 8): 2,
(64, 24, 8): 5,
(64, 32, 8): 2,
(64, 40, 8): 5,
(64, 48, 8): 5,
(64, 56, 8): 5,
(64, 64, 8): 1,
(64, 72, 8): 5,
(64, 80, 8): 5,
(64, 88, 8): 5,
(64, 96, 8): 5,
(64, 104, 8): 5,
(64, 112, 8): 5,
(64, 120, 8): 5,
(64, 128, 8): 0,
(64, 136, 8): 5,
(64, 144, 8): 5,
(64, 152, 8): 5,
(64, 160, 8): 5,
(64, 168, 8): 5,
(64, 176, 8): 5,
(64, 184, 8): 5,
(64, 192, 8): 3,
(64, 200, 8): 5,
(64, 208, 8): 5,
(64, 216, 8): 5,
(64, 224, 8): 5,
(64, 232, 8): 5,
(64, 240, 8): 5,
(64, 248, 8): 5,
(64, 256, 8): 3,
(128, 16, 8): 2,
(128, 32, 8): 2,
(128, 48, 8): 5,
(128, 64, 8): 2,
(128, 80, 8): 5,
(128, 96, 8): 5,
(128, 112, 8): 5,
(128, 128, 8): 0,
(128, 144, 8): 5,
(128, 160, 8): 5,
(128, 176, 8): 5,
(128, 192, 8): 3,
(128, 208, 8): 5,
(128, 224, 8): 5,
(128, 240, 8): 5,
(128, 256, 8): 0,
}
SM100_MMA_SHAPES_TF32_DENSE_2SM = {
(128, 32, 8): 2,
(128, 64, 8): 1,
(128, 96, 8): 5,
(128, 128, 8): 0,
(128, 160, 8): 5,
(128, 192, 8): 5,
(128, 224, 8): 5,
(128, 256, 8): 0,
(256, 32, 8): 2,
(256, 64, 8): 1,
(256, 96, 8): 5,
(256, 128, 8): 0,
(256, 160, 8): 5,
(256, 192, 8): 5,
(256, 224, 8): 5,
(256, 256, 8): 0,
}
# F8F6F4
SM100_MMA_SHAPES_F8F6F4_DENSE_1SM = {
(64, 8, 32): 4,
(64, 16, 32): 4,
(64, 24, 32): 5,
(64, 32, 32): 3,
(64, 40, 32): 5,
(64, 48, 32): 5,
(64, 56, 32): 5,
(64, 64, 32): 2,
(64, 72, 32): 5,
(64, 80, 32): 5,
(64, 88, 32): 5,
(64, 96, 32): 5,
(64, 104, 32): 5,
(64, 112, 32): 5,
(64, 120, 32): 5,
(64, 128, 32): 0,
(64, 136, 32): 5,
(64, 144, 32): 5,
(64, 152, 32): 5,
(64, 160, 32): 5,
(64, 168, 32): 5,
(64, 176, 32): 5,
(64, 184, 32): 5,
(64, 192, 32): 5,
(64, 200, 32): 5,
(64, 208, 32): 5,
(64, 216, 32): 5,
(64, 224, 32): 5,
(64, 232, 32): 5,
(64, 240, 32): 5,
(64, 248, 32): 5,
(64, 256, 32): 0,
(128, 16, 32): 4,
(128, 32, 32): 3,
(128, 48, 32): 5,
(128, 64, 32): 2,
(128, 80, 32): 5,
(128, 96, 32): 5,
(128, 112, 32): 5,
(128, 128, 32): 0,
(128, 144, 32): 5,
(128, 160, 32): 5,
(128, 176, 32): 5,
(128, 192, 32): 5,
(128, 208, 32): 5,
(128, 224, 32): 5,
(128, 240, 32): 5,
(128, 256, 32): 0,
}
SM100_MMA_SHAPES_F8F6F4_DENSE_2SM = {
(128, 32, 32): 3,
(128, 64, 32): 2,
(128, 96, 32): 5,
(128, 128, 32): 1,
(128, 160, 32): 5,
(128, 192, 32): 5,
(128, 224, 32): 5,
(128, 256, 32): 1,
(256, 32, 32): 2,
(256, 64, 32): 2,
(256, 96, 32): 5,
(256, 128, 32): 0,
(256, 160, 32): 5,
(256, 192, 32): 5,
(256, 224, 32): 5,
(256, 256, 32): 0,
}
# MXF8F6F4
SM100_MMA_SHAPES_MXF8F6F4_DENSE_1SM = {
(128, 64, 32): 1,
(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,
(256, 192, 32): 1,
(256, 256, 32): 0,
}
# MXF4NVF4
SM100_MMA_SHAPES_MXF4NVF4_DENSE_1SM = {
(128, 64, 64): 1,
(128, 128, 64): 0,
(128, 192, 64): 1,
(128, 256, 64): 0,
}
SM100_MMA_SHAPES_MXF4NVF4_DENSE_2SM = {
# Multiples of 16 for N
(256, 64, 64): 1,
(256, 128, 64): 0,
(256, 192, 64): 1,
(256, 256, 64): 0,
}
python/cutlass_library/sm100_utils.py
+661
View File
@@ -0,0 +1,661 @@
#################################################################################################
#
# Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#################################################################################################
"""
Utilities for enumerating CUTLASS library SM100 kernels
"""
import argparse
import enum
from itertools import product
import math
import logging
import os.path
import shutil
import sys
import copy
from typing import Any, Optional, Sequence, Tuple, List, Union, Callable
try:
import builtins
if hasattr(builtins, "CUTLASS_IGNORE_PACKAGE") and CUTLASS_IGNORE_PACKAGE == True:
raise ImportError("Disabling attempt to import cutlass_library")
from cutlass_library.library import *
except ImportError:
from library import *
#### Step 0: define levels
# One integer level controls multiple "generators" and how many
# combinations they generate. That is the "global" level.
# "Generators" are WGMMA shapes, MMA multipliers, cluster sizes, and
# anything that is eventually involved in the Cartesian product
# which yields our kernel configurations.
# For simplicity, each generator defines their own levels,
# starting from 0. As a rule we assume 10 or fewer levels, making
# their level a digit.
# The "global" level simply stacks these digits and represents them
# as a single integer.
#
# For example, level 500 indicates cluster sizes are at level 5, MMA
# multipliers are at level 0, and WGMMA shapes are at level 0 as well.
#
# Here we define the global level to generator level mappings.
def get_tcgen05_level_from_global_level(global_level: int):
return global_level % 10
def get_mma_level_from_global_level(global_level: int):
return (global_level // 10) % 10
def get_cluster_level_from_global_level(global_level: int):
return (global_level // 100) % 10
def get_pruning_level_from_global_level(global_level: int):
return (global_level // 1000) % 10
#### Step 1: generate MMA instruction shapes based on levels
try:
from .sm100_shapes import *
except:
from sm100_shapes import *
###########
def generate_tf32_math_instructions_sm100(level: int):
"""
Generate all TensorOp math instructions for TF32 MMA that are supported by SM100 at or above the given level.
Args:
level: The global level to generate math instructions for.
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_TF32_DENSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_TF32_DENSE_2SM.items() if tcgen05_level >= min_level
]
for shape in shapes_1sm:
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.tf32, DataType.tf32, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
for shape in shapes_2sm:
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.tf32, DataType.tf32, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
return math_instructions_1sm, math_instructions_2sm
def generate_16b_math_instructions_sm100(level: int):
"""
Generate all TensorOp math instructions for 16b MMA that are supported by SM100 at or above the given level.
Args:
level: The global level to generate math instructions for.
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_16b_DENSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_16b_DENSE_2SM.items() if tcgen05_level >= min_level
]
for shape in shapes_1sm:
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.f16, DataType.f16, DataType.f16,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.f16, DataType.f16, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.bf16, DataType.bf16, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
for shape in shapes_2sm:
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.f16, DataType.f16, DataType.f16,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.f16, DataType.f16, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.bf16, DataType.bf16, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
return math_instructions_1sm, math_instructions_2sm
def generate_fp8_math_instructions_sm100(level: int, enable_runtime_dtype = True, enable_compile_time_dtype = True):
"""
Generate all TensorOp math instructions for FP8 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_F8F6F4_DENSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_F8F6F4_DENSE_2SM.items() if tcgen05_level >= min_level
]
for shape in shapes_1sm:
if enable_runtime_dtype:
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.f8, DataType.f8, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
if enable_compile_time_dtype:
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.e4m3, DataType.e4m3, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.e5m2, DataType.e4m3, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.e4m3, DataType.e5m2, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
if pruning_level >= 2:
math_instructions_1sm.append(
MathInstruction(
shape,
DataType.e5m2, DataType.e5m2, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
for shape in shapes_2sm:
if enable_runtime_dtype:
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.f8, DataType.f8, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
if enable_compile_time_dtype:
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.e4m3, DataType.e4m3, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.e5m2, DataType.e4m3, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.e4m3, DataType.e5m2, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
if pruning_level >= 2:
math_instructions_2sm.append(
MathInstruction(
shape,
DataType.e5m2, DataType.e5m2, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
return math_instructions_1sm, math_instructions_2sm
def generate_f8f6f4_math_instructions_sm100(level: int, enable_runtime_dtype = True, enable_compile_time_dtype = True):
"""
Generate all TensorOp math instructions for FP8 FP6 and FP4 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_F8F6F4_DENSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_F8F6F4_DENSE_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):
math_instructions_1sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
if enable_compile_time_dtype:
compile_time_types = [ DataType.e4m3, DataType.e5m2, DataType.e3m2, 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.TensorOp,
MathOperation.multiply_add)
)
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):
math_instructions_2sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.TensorOp,
MathOperation.multiply_add)
)
if enable_compile_time_dtype:
compile_time_types = [ DataType.e4m3, DataType.e5m2, DataType.e3m2, 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.TensorOp,
MathOperation.multiply_add)
)
return math_instructions_1sm, math_instructions_2sm
def generate_mxf8f6f4_math_instructions_sm100(level: int, enable_runtime_dtype = True, enable_compile_time_dtype = True):
"""
Generate all BlockScaledTensorOp 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_DENSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_MXF8F6F4_DENSE_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.BlockScaledTensorOp,
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.BlockScaledTensorOp,
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.BlockScaledTensorOp,
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.BlockScaledTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
return math_instructions_1sm, math_instructions_2sm
def generate_mxf4nvf4_math_instructions_sm100(level: int, enable_runtime_dtype = True, enable_compile_time_dtype = True):
"""
Generate all BlockScaledTensorOp 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_DENSE_1SM.items() if tcgen05_level >= min_level
]
shapes_2sm = [
shape for shape, min_level in SM100_MMA_SHAPES_MXF4NVF4_DENSE_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.BlockScaledTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
math_instructions_1sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledTensorOp,
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.BlockScaledTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
math_instructions_1sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledTensorOp,
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.BlockScaledTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
math_instructions_2sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledTensorOp,
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.BlockScaledTensorOp,
MathOperation.multiply_add,
DataType.ue8m0)
)
math_instructions_2sm.append(
MathInstruction(
shape,
a_type, b_type, DataType.f32,
OpcodeClass.BlockScaledTensorOp,
MathOperation.multiply_add,
DataType.ue4m3)
)
return math_instructions_1sm, math_instructions_2sm
def generate_cluster_shapes_sm100(level: int, change_priority_func : Union[Callable, None] = None):
"""
Generate all cluster shapes for SM100 at or above the given level.
Args:
level: The global level to generate cluster shapes for.
Returns:
A tuple of two lists of cluster shapes.
The first list contains the cluster shapes for 1SM, and the second list contains the cluster shapes for 2SM.
"""
cluster_level = get_cluster_level_from_global_level(level)
assert cluster_level >= 4
if change_priority_func is not None:
SM100_CLUSTER_SHAPES_1SM_CPY = copy.deepcopy(SM100_CLUSTER_SHAPES_1SM)
SM100_CLUSTER_SHAPES_2SM_CPY = copy.deepcopy(SM100_CLUSTER_SHAPES_2SM)
change_priority_func(SM100_CLUSTER_SHAPES_1SM_CPY, SM100_CLUSTER_SHAPES_2SM_CPY)
shapes_1sm = [
list(shape) for shape, min_level in SM100_CLUSTER_SHAPES_1SM_CPY.items() if cluster_level >= min_level
]
shapes_2sm = [
list(shape) for shape, min_level in SM100_CLUSTER_SHAPES_2SM_CPY.items() if cluster_level >= min_level
]
return shapes_1sm, shapes_2sm
else:
shapes_1sm = [
list(shape) for shape, min_level in SM100_CLUSTER_SHAPES_1SM.items() if cluster_level >= min_level
]
shapes_2sm = [
list(shape) for shape, min_level in SM100_CLUSTER_SHAPES_2SM.items() if cluster_level >= min_level
]
return shapes_1sm, shapes_2sm
python/cutlass_library/sm90_utils.py
+4 -1
View File
@@ -637,7 +637,10 @@ def get_valid_schedules(tile_description, cuda_version, is_aligned, data_types,
if CudaToolkitVersionSatisfies(cuda_version, 12, 1):
# Pruning: don't stamp out fp8 ping-pong kernel with non-tma epilogue
if not is_fp8 or level >= 1:
schedules.append([to_grouped_schedule(KernelScheduleType.TmaWarpSpecializedPingpong, grouped), to_grouped_schedule(default_epilogue, grouped)])
if not is_blockwise(gemm_kind):
schedules.append([to_grouped_schedule(KernelScheduleType.TmaWarpSpecializedPingpong, grouped), to_grouped_schedule(default_epilogue, grouped)])
else:
schedules.append([to_grouped_schedule(KernelScheduleType.BlockwiseTmaWarpSpecializedPingpong, grouped), to_grouped_schedule(default_epilogue, grouped)])
if can_do_fp8_fast_accum:
if not grouped: