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

build_configs

def build_configs[a_type: DType, b_type: DType, c_type: DType, N: Int, K: Int, transpose_b: Bool = True, num_k_partitions: Int = Int(1), partitioned_multicast: Bool = False, pdl_level: PDLLevel = PDLLevel.OFF, k_groups: Optional[Int] = None, consumer_groups: Optional[Int] = None, swapAB: Bool = False]() -> Set[MatmulConfig[a_type, b_type, c_type, transpose_b]]

Builds the set of unique SM90 MatmulConfig instances for a fixed (N, K) and all M values.

Sweeps M from 8 to 8192 to enumerate the distinct configs the heuristic produces for the given (N, K). Duplicate configs (same tile/cluster shape) are deduplicated via set membership.

Parameters:

a_type (DType): Element DType of the A input matrix.
b_type (DType): Element DType of the B input matrix.
c_type (DType): Element DType of the output C matrix.
N (Int): The N dimension (columns) of the matmul, fixed across the M sweep.
K (Int): The K dimension (contraction axis) of the matmul.
transpose_b (Bool): Whether B is stored transposed (K-major layout) (defaults to True).
num_k_partitions (Int): Number of partitions to split the K dimension into (defaults to 1).
partitioned_multicast (Bool): Whether to use partitioned TMA multicast (defaults to False).
pdl_level (PDLLevel): Programmatic dependent launch level for grid controls (defaults to PDLLevel.OFF).
k_groups (Optional[Int]): Optional override for the K-grouping size; None lets the heuristic decide (defaults to None).
consumer_groups (Optional[Int]): Optional override for the number of consumer warp groups; None lets the heuristic decide (defaults to None).
swapAB (Bool): Whether to swap the A and B roles, treating N as M (defaults to False).

Returns:

Set[MatmulConfig[a_type, b_type, c_type, transpose_b]]: A set of unique MatmulConfig values covering the M sweep.