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

swapAB_smallM

def swapAB_smallM[a_type: DType, b_type: DType, c_type: DType, prioritize_compute_over_ctas: Bool = False, transpose_b: Bool = True](m: Int, n: Int, k: Int, cluster_shape: IndexList[Int(3)], num_k_partitions: Int, num_consumer: Int, partitioned_multicast: Bool, pdl_level: PDLLevel, k_group_size: Int = Int(0), num_pipeline_stages: Int = Int(0)) -> MatmulConfig[a_type, b_type, c_type, transpose_b]

Selects an SM90 MatmulConfig for small-M problems by swapping A and B roles.

When M is small compared to N, treating N as the leading dimension and swapping A/B improves SM utilization. Sweeps MMA-N values and picks the configuration that maximizes SM wave occupancy or compute ratio.

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.
prioritize_compute_over_ctas (Bool): Whether to select MMA-N by minimizing wasted compute ratio rather than maximizing CTAs used (defaults to False).
transpose_b (Bool): Whether B is stored transposed (K-major layout) (defaults to True).

Args:

m (Int): The M dimension of the matmul (mapped to N after the A/B swap).
n (Int): The N dimension of the matmul (mapped to M after the A/B swap).
k (Int): The K dimension (contraction axis) of the matmul.
cluster_shape (IndexList[Int(3)]): The thread block cluster dimensions in (M, N, K).
num_k_partitions (Int): Number of partitions to split the K dimension into.
num_consumer (Int): Number of consumer warp groups.
partitioned_multicast (Bool): Whether to use partitioned TMA multicast.
pdl_level (PDLLevel): Programmatic dependent launch level for grid controls.
k_group_size (Int): Number of K iterations grouped per pipeline stage, or 0 to auto-compute (defaults to 0).
num_pipeline_stages (Int): Number of pipeline stages for double-buffering loads, or 0 to maximize by default (defaults to 0).

Returns:

MatmulConfig[a_type, b_type, c_type, transpose_b]: A MatmulConfig with swapped AB dimensions tuned for small-M inputs.