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_tile_rank: Int, a_tile_shape: IndexList[a_tile_rank], a_desc_shape: IndexList[a_tile_rank], b_tile_rank: Int, b_tile_shape: IndexList[b_tile_rank], b_desc_shape: IndexList[b_tile_rank], c_tile_rank: Int, c_tile_shape_param: IndexList[c_tile_rank], c_tensor_layout: Layout, c_desc_shape: IndexList[c_tile_rank], sfa_tile_rank: Int, sfa_tile_shape: IndexList[sfa_tile_rank], sfa_desc_shape: IndexList[sfa_tile_rank], sfb_tile_rank: Int, sfb_tile_shape: IndexList[sfb_tile_rank], sfb_desc_shape: IndexList[sfb_tile_rank], group_offsets_layout: Layout, group_scale_offsets_layout: Layout, expert_ids_layout: Layout, expert_scales_layout: Layout, transpose_b: Bool, config: BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b], expert_n: Int, 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, pdl_level: PDLLevel = PDLLevel(1), max_profiled_tiles_per_SM: UInt32 = UInt32(0)](num_active_experts: Int, a_tma_op: TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape], b_tma_op: TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape], c_tma_op: TMATensorTile[c_type, c_tile_rank, c_tile_shape_param, c_desc_shape], c: LayoutTensor[c_type, c_tensor_layout, MutAnyOrigin], sfa_tma_op: TMATensorTile[sfa_dtype, sfa_tile_rank, sfa_tile_shape, sfa_desc_shape], sfb_tma_op: TMATensorTile[sfb_dtype, sfb_tile_rank, sfb_tile_shape, sfb_desc_shape], group_offsets: LayoutTensor[DType.uint32, group_offsets_layout, MutAnyOrigin], group_scale_offsets: LayoutTensor[DType.uint32, group_scale_offsets_layout, MutAnyOrigin], expert_ids: LayoutTensor[DType.int32, expert_ids_layout, MutAnyOrigin], expert_scales: LayoutTensor[DType.float32, expert_scales_layout, MutAnyOrigin], cluster_dim: StaticTuple[Int32, 3], mnk: StaticTuple[UInt32, 3], workspace: Span[UInt64, MutAnyOrigin])