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

swiglu_matmul_config

def swiglu_matmul_config[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool = True](*, mma_shape: IndexList[Int(3)], cta_group: Int = Int(2), cluster_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(2), Int(1), Int(1)), AB_swapped: Bool = False, block_swizzle_size: Int = Int(0), raster_order: RasterOrder = RasterOrder.AlongM, k_group_size: Int = Int(1), num_clc_pipeline_stages: Int = Int(2), num_accum_pipeline_stages: Int = Int(2), use_bias: Bool = False, register_swiglu: Bool = False) -> FusedSwiGLUMatmulConfig[a_type, b_type, c_type, transpose_b]

Create a FusedSwiGLUMatmulConfig with auto-computed pipeline stages.

For the SMEM→TMA path (register_swiglu=False) the extra SMEM consumed by the double-buffered half-output tiles is subtracted from the budget before solving for the stage count. For the register→GMEM path (register_swiglu=True) no extra deduction is made.

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).

Args:

mma_shape (IndexList[Int(3)]): MMA instruction shape as (MMA_M, MMA_N, MMA_K).
cta_group (Int): Number of cooperating CTAs per cluster group (defaults to 2).
cluster_shape (IndexList[Int(3)]): Thread block cluster tile counts as (cluster_m, cluster_n, cluster_k) (defaults to (2, 1, 1)).
AB_swapped (Bool): Whether the A and B operands are swapped (defaults to False).
block_swizzle_size (Int): Block swizzle factor for tile remapping, one of 0, 1, 2, 4, or 8 (defaults to 0).
raster_order (RasterOrder): Order CLC rasterizes tiles across the cluster grid (defaults to RasterOrder.AlongM).
k_group_size (Int): Number of K tiles loaded per pipeline stage (defaults to 1).
num_clc_pipeline_stages (Int): Number of CLC pipeline stages for work distribution (defaults to 2).
num_accum_pipeline_stages (Int): Number of accumulator pipeline stages (defaults to 2).
use_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 False).

Returns:

FusedSwiGLUMatmulConfig[a_type, b_type, c_type, transpose_b]