Mojo struct
MatmulKernels
struct MatmulKernels[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool = False]
Supported matmul kernels.
The configurations are named as:
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
ampere_128x128_4
comptime ampere_128x128_4 = MatmulConfig(Index(128, 128, _bk_base[a_type]()), Index(64, 64, _bk_base[a_type]()), get_mma_shape[a_type, MatmulConfig[a_type, b_type, c_type, transpose_b].accum_type](), Index(1, 1, 1), 4, 1, 1, 1, 1, False, PDLLevel())
ampere_256x128_3
comptime ampere_256x128_3 = MatmulConfig(Index(128, 256, (2 * _bk_base[a_type]())), Index(64, 64, (2 * _bk_base[a_type]())), get_mma_shape[a_type, MatmulConfig[a_type, b_type, c_type, transpose_b].accum_type](), Index(1, 1, 1), 3, 1, 1, 1, 1, False, PDLLevel())
ampere_256x64_4
comptime ampere_256x64_4 = MatmulConfig(Index(64, 256, _bk_base[a_type]()), Index(64, 64, _bk_base[a_type]()), get_mma_shape[a_type, MatmulConfig[a_type, b_type, c_type, transpose_b].accum_type](), Index(1, 1, 1), 4, 1, 1, 1, 1, False, PDLLevel())
hopper_128x128_4
comptime hopper_128x128_4 = MatmulConfig(Index(128, 128, _bk_base[a_type]()), Index(64, 64, _bk_base[a_type]()), get_mma_shape[a_type, MatmulConfig[a_type, b_type, c_type, transpose_b].accum_type](), Index(1, 1, 1), 4, 1, 1, 1, 1, False, PDLLevel())
tuning_config
comptime tuning_config = MatmulConfig(Index(get_defined_int["TUNE_BM", 128](), get_defined_int["TUNE_BN", 128](), get_defined_int["TUNE_BK", 32]()), Index(get_defined_int["TUNE_WM", 64](), get_defined_int["TUNE_WN", 64](), get_defined_int["TUNE_BK", 32]()), get_mma_shape[a_type, MatmulConfig[a_type, b_type, c_type, transpose_b].accum_type](), Index(1, 1, 1), SIMD(get_defined_int["TUNE_NUM_STAGES", 4]()), SIMD(get_defined_int["TUNE_NUM_K_PARTITIONS", 1]()), 1, SIMD(get_defined_int["TUNE_NUM_WARP_K_PARTITIONS", 1]()), 1, False, PDLLevel())
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!