For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

choose_block_scaled_config

def choose_block_scaled_config[a_type: DType, b_type: DType, c_type: DType, sfa_dtype: DType, sfb_dtype: DType, transpose_b: Bool = True, gemm_kind: GEMMKind = GEMMKind.GEMM](M: Int, N: Int, K: Int) -> BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b]

Select a BlockScaledMatmulConfig that minimizes waves per SM for the given shape.

Parameters:

• ​a_type (DType): DType of the A (left) operand elements; uint8 indicates packed FP4; must equal b_type.
• ​b_type (DType): DType of the B (right) operand elements.
• ​c_type (DType): DType of the output matrix elements.
• ​sfa_dtype (DType): DType of the A operand block scaling factors; must equal sfb_dtype.
• ​sfb_dtype (DType): DType of the B operand block scaling factors; must equal sfa_dtype.
• ​transpose_b (Bool): Whether the B operand is stored transposed (defaults to True).
• ​gemm_kind (GEMMKind): The GEMMKind selecting the kernel variant (defaults to GEMMKind.GEMM).

Args:

• ​M (Int): The M dimension of the matmul.
• ​N (Int): The N dimension of the matmul.
• ​K (Int): The K dimension of the matmul.

Returns:

BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b]: A BlockScaledMatmulConfig tuned for the given problem dimensions.