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_sm100

def grouped_matmul_sm100[c_type: DType, a_type: DType, b_type: DType, //, *, transpose_b: Bool = True, mma_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(64), Int(128), Int(16)), block_tile_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(64), 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 SM100 grouped matmul kernel with TMA descriptors for ragged MoE matrix multiplication on Blackwell GPUs.

Was this page helpful?