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_blockwise_scaled_fp8

def grouped_matmul_sm100_blockwise_scaled_fp8[c_type: DType, a_type: DType, b_type: DType, a_scales_type: DType, b_scales_type: DType, a_offsets_type: DType, expert_ids_type: DType, transpose_b: Bool, //, *, config: MatmulConfig[a_type, b_type, c_type, transpose_b], 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], b: TileTensor[b_type, Storage=b.Storage, linear_idx_type=b.linear_idx_type], a_scales: TileTensor[a_scales_type, Storage=a_scales.Storage, linear_idx_type=a_scales.linear_idx_type], b_scales: TileTensor[b_scales_type, Storage=b_scales.Storage, linear_idx_type=b_scales.linear_idx_type], a_offsets: TileTensor[a_offsets_type, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type], expert_ids: TileTensor[expert_ids_type, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type], max_num_tokens_per_expert: Int, num_active_experts: Int, ctx: DeviceContext)

Launches the basic (non-persistent) SM100 blockwise-scaled FP8 grouped GEMM kernel.

Converts the input TileTensors to LayoutTensors, builds TMA descriptors for A, B, and the C output, and enqueues matmul_sm100_grouped_blockwise_scaled_fp8_1d2d_kernel with a grid of (N/BN, max_tokens/BM, num_active_experts) blocks.

Parameters:

c_type (DType): Element type of the output C tensor (inferred).
a_type (DType): Element type of the input A tensor (inferred). Must be float8_e4m3fn.
b_type (DType): Element type of the input B tensor (inferred). Must be float8_e4m3fn.
a_scales_type (DType): Element type of the a_scales tensor (inferred).
b_scales_type (DType): Element type of the b_scales tensor (inferred).
a_offsets_type (DType): Element type of the a_offsets tensor (inferred).
expert_ids_type (DType): Element type of the expert_ids tensor (inferred).
transpose_b (Bool): Whether B is stored transposed (inferred). Must be True.
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Matmul configuration specifying block tile shape, MMA shape, and TMA swizzle modes.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue function applied to each output element before storing (defaults to None).

Args:

c (TileTensor[c_type, Storage=c.Storage, linear_idx_type=c.linear_idx_type]): Output tensor of shape [total_tokens, N] holding the grouped matmul results.
a (TileTensor[a_type, Storage=a.Storage, linear_idx_type=a.linear_idx_type]): Input activation tensor of shape [total_tokens, K] in FP8.
b (TileTensor[b_type, Storage=b.Storage, linear_idx_type=b.linear_idx_type]): Input weight tensor of shape [num_experts, N, K] in FP8.
a_scales (TileTensor[a_scales_type, Storage=a_scales.Storage, linear_idx_type=a_scales.linear_idx_type]): Per-block scales for A of shape [K // BK, total_tokens] where BK is the scaling block size.
b_scales (TileTensor[b_scales_type, Storage=b_scales.Storage, linear_idx_type=b_scales.linear_idx_type]): Per-block scales for B of shape [num_experts, N // BN, K // BK] where BN and BK are the scaling block sizes.
a_offsets (TileTensor[a_offsets_type, Storage=a_offsets.Storage, linear_idx_type=a_offsets.linear_idx_type]): Cumulative row offsets per expert, length num_active_experts + 1. Entry i + 1 minus entry i gives the row count for expert i.
expert_ids (TileTensor[expert_ids_type, Storage=expert_ids.Storage, linear_idx_type=expert_ids.linear_idx_type]): Expert index for each active expert slot, mapping the grid Z index to the corresponding row offset in B.
max_num_tokens_per_expert (Int): Maximum number of tokens assigned to any single expert, used to size the grid M dimension.
num_active_experts (Int): Number of active experts in this grouped matmul, used to size the grid Z dimension.
ctx (DeviceContext): Device context used to enqueue the kernel.