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

enqueue_apple_matmul_split_k

def enqueue_apple_matmul_split_k[in_type: DType, c_type: DType = DType.float32, transpose_b: Bool = False, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[in_type, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[in_type, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type], ctx: DeviceContext, num_splits_hint: Int = Int(4))

Split-K Apple M5 matmul: partition K, accumulate partials, reduce.

num_splits_hint is an upper bound; the actual split count is capped so no split is empty (actual_splits = ceil(num_strips / strips_per_split) where strips_per_split = ceil(num_strips / num_splits_hint)). Best for large-K, small-M*N shapes where the single-pass kernel under-occupies the GPU.

Parameters:

in_type (DType): A/B element type (fp16, bf16, fp32).
c_type (DType): Output element type (fp16, bf16, fp32). Accumulation is fp32 (defaults to float32).
transpose_b (Bool): If True, B is (N, K) row-major (viewed as col_major(K, N)); otherwise B is (K, N) row-major (defaults to False).
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional fused epilogue; receives SIMD[c_type, width] at absolute (row, col) (AMD's contract) (defaults to None).

Args:

c (TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output matrix (M, N) row-major; M and N derive from its dims.
a (TileTensor[in_type, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): A operand matrix (M, K) row-major; K derives from a.dim[1].
b (TileTensor[in_type, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): B operand matrix, (K, N) for transpose_b=False or (N, K) for transpose_b=True.
ctx (DeviceContext): Device context to enqueue the kernels on; must be Apple M5 (compute_capability == 5).
num_splits_hint (Int): Upper bound on the K-axis split count; the actual split count is capped so no split is empty (defaults to 4).

Raises:

If the attached GPU is not Apple M5 (compute_capability != 5).