Skip to main content

Mojo function

blackwell_block_scaled_tma_umma_warp_specialized_kernel

blackwell_block_scaled_tma_umma_warp_specialized_kernel[a_type: DType, b_type: DType, c_type: DType, sfa_dtype: DType, sfb_dtype: DType, a_rank: Int, a_tile_shape: IndexList[a_rank], a_desc_shape: IndexList[a_rank], b_rank: Int, b_tile_shape: IndexList[b_rank], b_desc_shape: IndexList[b_rank], c_rank: Int, c_tile_shape: IndexList[c_rank], c_desc_shape: IndexList[c_rank], sfa_tma_dtype: DType, sfa_rank: Int, sfa_tile_shape: IndexList[sfa_rank], sfa_desc_shape: IndexList[sfa_rank], transpose_b: Bool, config: BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b], cluster_shape: StaticTuple[Int32, 3] = StaticTuple(Int32(1)), elementwise_compute_lambda_fn: Optional[def[dtype: DType, width: Int, *, alignment: Int = 1](IndexList[2], SIMD[dtype, width]) capturing -> SIMD[dtype, width]] = None, elementwise_lambda_fn: Optional[def[dtype: DType, width: Int, *, alignment: Int = 1](IndexList[2], SIMD[dtype, width]) capturing -> None] = None, pdl_level: PDLLevel = PDLLevel(), max_profiled_tiles_per_SM: UInt32 = UInt32(0)](a_tma_op: TMATensorTile[a_type, a_rank, a_tile_shape, a_desc_shape], b_tma_op: TMATensorTile[b_type, b_rank, b_tile_shape, b_desc_shape], c_tma_op: TMATensorTile[c_type, c_rank, c_tile_shape, c_desc_shape], sfa_tma_op: TMATensorTile[sfa_tma_dtype, sfa_rank, sfa_tile_shape, sfa_desc_shape], cluster_dim: StaticTuple[Int32, 3], mnk: StaticTuple[UInt32, 3], workspace: Span[UInt64, MutAnyOrigin], alpha: Float32, sfb_global_ptr: UnsafePointer[Scalar[sfb_dtype], ImmutAnyOrigin], sfb_batch_stride: Int, sfb_n_stride: Int, sfb_k_tiles: Int)

Was this page helpful?