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

choose_swiglu_config

def choose_swiglu_config[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool = True, has_bias: Bool = False, register_swiglu: Bool = True](M: Int, N: Int, K: Int) -> FusedSwiGLUMatmulConfig[a_type, b_type, c_type, transpose_b]

Select a FusedSwiGLUMatmulConfig using a wave-minimization heuristic.

Implements the same MMA shape selection as the normal matmul heuristic, simplified for BF16 (no FP8 special cases). For small M (< 32) uses a single-CTA configuration with AB swapped; for large M uses two-CTA clusters and sweeps both normal and swapped AB to minimize waves. FusedSwiGLUMatmulConfig.__init__ recomputes num_pipeline_stages from the swiglu-specific SMEM budget (which accounts for the half-output tile and optional bias tile overhead).

Parameters:

a_type (DType): DType of the A operand elements.
b_type (DType): DType of the B operand elements.
c_type (DType): DType of the C output elements.
transpose_b (Bool): Whether the B operand is stored transposed (defaults to True).
has_bias (Bool): Whether a bias vector is applied in the epilogue (defaults to False).
register_swiglu (Bool): Selects the SwiGLU epilogue path: True for register→GMEM, False for SMEM→TMA (defaults to True).

Args:

M (Int): Number of rows in the output C matrix.
N (Int): Number of columns in the output C matrix.
K (Int): Contraction dimension of the GEMM (shared inner size of A and B).

Returns:

FusedSwiGLUMatmulConfig[a_type, b_type, c_type, transpose_b]