For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

AppleM5Int8MatMul

struct AppleM5Int8MatMul[c_type: DType = DType.bfloat16, *, has_bias: Bool = False, BM: Int = Int(64), BN: Int = Int(64), BK: Int = Int(64), TTI32: Bool = False]

W8A8 GEMM: int8 A x int8 B^T -> int32 accum -> per-row/col dequant.

run is the GPU kernel entry (TileTensor operands + M/N/K). Launch via enqueue_apple_int8_matmul.

Parameters

• ​c_type (DType): Output element type (bf16 / fp16 / fp32). Accumulation is int32; the dequant multiply is done in fp32 then cast to c_type.
• ​has_bias (Bool): If True, add a per-output-column bias (in c_type) after dequant.
• ​BM (Int): Threadgroup M-tile height (multiple of SG_M).
• ​BN (Int): Threadgroup N-tile width (multiple of SG_N).
• ​BK (Int): K-strip depth per accumulate step (multiple of MMA_K = 16; default 64 = four 16-wide K-blocks).
• ​TTI32 (Bool): If True, use int32 load_linear for the interior A/B loads (faster on NA-bound shapes, numerically identical). Auto-selected by enqueue_apple_int8_matmul when max(M*K, N*K) < 2^31, else i64.

Implemented traits

AnyType, Deinitable, Movable

comptime members

Mma

comptime Mma = MmaOpApple[DType.int32, DType.int8, Int(2), Int(2), transpose_b=True]

MMA_K

comptime MMA_K = 16

MMA_M

comptime MMA_M = 16

MMA_N

comptime MMA_N = 16

NUM_MMA_M

comptime NUM_MMA_M = 2

NUM_MMA_N

comptime NUM_MMA_N = 2

NUM_SG

comptime NUM_SG = (AppleM5Int8MatMul[c_type, has_bias=has_bias, BM=BM, BN=BN, BK=BK, TTI32=TTI32].NUM_SG_M * AppleM5Int8MatMul[c_type, has_bias=has_bias, BM=BM, BN=BN, BK=BK, TTI32=TTI32].NUM_SG_N)

NUM_SG_M

comptime NUM_SG_M = (BM // Int(32))

NUM_SG_N

comptime NUM_SG_N = (BN // Int(32))

SG_M

comptime SG_M = 32

SG_N

comptime SG_N = 32

THREADS_PER_BLOCK

comptime THREADS_PER_BLOCK = (AppleM5Int8MatMul[c_type, has_bias=has_bias, BM=BM, BN=BN, BK=BK, TTI32=TTI32].NUM_SG * _resolve_warp_size())

Methods

morton_decode_2d

static def morton_decode_2d(flat_idx: UInt32) -> Tuple[UInt32, UInt32]

Returns:

Tuple[UInt32, UInt32]

morton_decode_2d_rect

static def morton_decode_2d_rect(flat_idx: UInt32, log2_m: UInt32, log2_n: UInt32) -> Tuple[UInt32, UInt32]

Returns:

Tuple[UInt32, UInt32]

run

static def run[c_layout: TensorLayout, a_layout: TensorLayout, b_layout: TensorLayout, as_layout: TensorLayout, bs_layout: TensorLayout, bias_layout: TensorLayout](c: TileTensor[c_type, c_layout, MutAnyOrigin], a: TileTensor[DType.int8, a_layout, ImmutAnyOrigin], b: TileTensor[DType.int8, b_layout, ImmutAnyOrigin], a_scale: TileTensor[DType.float32, as_layout, ImmutAnyOrigin], b_scale: TileTensor[DType.float32, bs_layout, ImmutAnyOrigin], bias: TileTensor[c_type, bias_layout, ImmutAnyOrigin], log2_grid_m: UInt32, log2_grid_n: UInt32)

W8A8 kernel entry. C (M, N), A (M, K) int8, B (N, K) int8 (transpose_b), a_scale (M,), b_scale (N,), bias (N,) (used iff has_bias). Grid is (1<<log2_grid_m) * (1<<log2_grid_n) threadgroups; OOB threadgroups early-return after Morton decode.

Parameters:

• ​c_layout (TensorLayout): TensorLayout of the C output TileTensor.
• ​a_layout (TensorLayout): TensorLayout of the int8 A TileTensor.
• ​b_layout (TensorLayout): TensorLayout of the int8 B TileTensor.
• ​as_layout (TensorLayout): TensorLayout of the per-row activation scale TileTensor.
• ​bs_layout (TensorLayout): TensorLayout of the per-column weight scale TileTensor.
• ​bias_layout (TensorLayout): TensorLayout of the per-column bias TileTensor.

Args:

• ​c (TileTensor[c_type, c_layout, MutAnyOrigin]): Output TileTensor of shape (M, N) in c_type.
• ​a (TileTensor[DType.int8, a_layout, ImmutAnyOrigin]): Int8 activation TileTensor of shape (M, K).
• ​b (TileTensor[DType.int8, b_layout, ImmutAnyOrigin]): Int8 weight TileTensor of shape (N, K) (used transposed).
• ​a_scale (TileTensor[DType.float32, as_layout, ImmutAnyOrigin]): Per-row fp32 activation scale, shape (M,).
• ​b_scale (TileTensor[DType.float32, bs_layout, ImmutAnyOrigin]): Per-column fp32 weight scale, shape (N,).
• ​bias (TileTensor[c_type, bias_layout, ImmutAnyOrigin]): Per-column bias in c_type, shape (N,) (used iff has_bias).
• ​log2_grid_m (UInt32): Base-2 log of the power-of-two-padded M-side grid extent.
• ​log2_grid_n (UInt32): Base-2 log of the power-of-two-padded N-side grid extent.