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 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, ImplicitlyDeletable

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:

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact