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_layout: Layout, b_layout: Layout, c_layout: Layout, c_tensor_layout: Layout, sfa_tile_layout: Layout, sfb_tile_layout: Layout, a_desc_layout: Layout, b_desc_layout: Layout, c_desc_layout: Layout, sfa_desc_layout: Layout, sfb_desc_layout: Layout, 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(1), elementwise_compute_lambda_fn: Optional[elementwise_compute_lambda_type] = None, register_based_epilogue: Bool = True, pdl_level: PDLLevel = PDLLevel(), max_profiled_tiles_per_SM: UInt32 = 0](num_active_experts: Int, a_tma_op: TMATensorTile[a_type, a_layout, a_desc_layout], b_tma_op: TMATensorTile[b_type, b_layout, b_desc_layout], c_tma_op: TMATensorTile[c_type, c_layout, c_desc_layout], c: LayoutTensor[c_type, c_tensor_layout, MutAnyOrigin], sfa_tma_op: TMATensorTile[sfa_dtype, sfa_tile_layout, sfa_desc_layout], sfb_tma_op: TMATensorTile[sfb_dtype, sfb_tile_layout, sfb_desc_layout], 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])
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!