Skip to main content

Mojo function

test_blackwell_block_scaled_matmul_tma_umma_warp_specialized

test_blackwell_block_scaled_matmul_tma_umma_warp_specialized[MType: CoordLike, NType: CoordLike, KType: CoordLike, //, a_type: DType, b_type: DType, c_type: DType, scales_dtype: DType, block_tile_shape: IndexList[3], mma_shape: IndexList[3], cluster_shape: StaticTuple[Int32, 3], cta_group: Int, transpose_b: Bool = True, a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, c_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, block_swizzle_size: Int = 0, benchmark: Bool = False, swapAB: Bool = False, k_group_size: Int = 1, num_clc_pipeline_stages: Int = 2, SF_VECTOR_SIZE: Int = 16, is_small_bn: Bool = False, normal_epilogue: Bool = False](ctx: DeviceContext, m: MType, n: NType, k: KType, alpha: Float32 = 1)

Was this page helpful?