IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

small_MN_gemms

def small_MN_gemms[config: TuningConfigSmallMNGemms, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, pdl_level: PDLLevel = PDLLevel()](c: TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type], ctx: DeviceContext)

Launches a small-MN GEMM via the configured split-K GEMV or MMA-CPasync kernel.

Selects between gemm_mma_cpasync (for GEMM_MMA_CPASYNC kernel kind) and gemv_split_k (otherwise) based on config.kernel_kind, then enqueues the chosen kernel with the runtime M, N, K derived from the input tiles.

Parameters:

Was this page helpful?