IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

grouped_matmul_amd

def grouped_matmul_amd[c_type: DType, a_type: DType, b_type: DType, *, transpose_b: Bool = True, block_tile_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(128), Int(128), Int(64)), elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, Storage=c.Storage, linear_idx_type=c.linear_idx_type], a: TileTensor[a_type, Storage=a.Storage, linear_idx_type=a.linear_idx_type], a_offsets: TileTensor[DType.uint32, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type], max_num_tokens_per_expert: Int, b: TileTensor[b_type, Storage=b.Storage, linear_idx_type=b.linear_idx_type], expert_ids: TileTensor[DType.int32, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type], num_active_experts: Int, ctx: DeviceContext)

Launches the AMD grouped matmul kernel, selecting the best block-tile configuration for the runtime M dimension via dispatch_amd_matmul_by_block_shape.

Was this page helpful?