Skip to main content

Mojo function

blackwell_gmm_tma_umma_warp_specialized_blockwise_fp8_kernel

blackwell_gmm_tma_umma_warp_specialized_blockwise_fp8_kernel[a_type: DType, b_type: DType, c_type: DType, a_layout: Layout, b_layout: Layout, c_layout: Layout, c_tensor_layout: Layout, a_scales_tile_layout: Layout, a_scales_type: DType, a_offsets_layout: Layout, b_scales_type: DType, b_scales_layout: Layout, a_desc_layout: Layout, b_desc_layout: Layout, c_desc_layout: Layout, a_scales_desc_layout: Layout, transpose_b: Bool, config: MatmulConfig[a_type, b_type, c_type, transpose_b], num_pipeline_stages: UInt, cluster_shape: StaticTuple[Int32, 3], expert_n: Int, expert_ids_layout: Layout, b_scales_n: Int](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], a_scales_tma_op: TMATensorTile[a_scales_type, a_scales_tile_layout, a_scales_desc_layout], a_offsets: LayoutTensor[DType.uint32, a_offsets_layout, MutAnyOrigin], num_iters: UInt, b_scales: LayoutTensor[b_scales_type, b_scales_layout, MutAnyOrigin], expert_ids: LayoutTensor[DType.int32, expert_ids_layout, MutAnyOrigin], problem_shape: StaticTuple[Int32, 3])

Was this page helpful?