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_fp8_matmul

def enqueue_apple_fp8_matmul[c_type: DType = DType.float32, 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[DType.bfloat16, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], weight: TileTensor[DType.float8_e4m3fn, Storage=weight.Storage, address_space=weight.address_space, linear_idx_type=weight.linear_idx_type], ctx: DeviceContext)

Enqueue the W8A16 matmul: out = a @ W_fp8^T (raw, unscaled).

a is the bf16 activation (M, K), weight the FP8-E4M3 weight (N, K) (transpose_b). C is (M, N). Produces the RAW matmul; the per-tensor scalar weight_scale is folded post-matmul by the graph lowering (quant_ops._matmul_float8, the FP8 analog of NVFP4's weight_scale_2).

Dispatch (all routes produce consistent weight values -- E4M3 -> f32/bf16 is exact):

Tiled FP8 matmul (enqueue_matmul2d_fp8; bf16 A x fp8 B -> fp32 on the native Apple MMA, direct DRAM fp8 feed), no fused epilogue, for the wide-N decode class (n > k: Mamba in_proj, MLP up) at ANY M AND every M > 1 shape including narrow-N (out_proj, MLP down). It amortizes the per-output cost across a threadgroup tile, beating both the bf16 matmul and the materialize route (measured M5 Max, M=32: 1.20-2.17x the bf16 matmul on all four Nemotron FP8 Linears; the prior narrow-N M>1 materialize route ran at 0.11-0.20x bf16 -- the c32 FP8 regression).
Batch-1 decode (M == 1) narrow-N (n <= k: out_proj, MLP down): the register-resident W8A16 GEMV (enqueue_apple_fp8_gemv), no MMA. Its long per-warp K amortizes the per-output cost, so it wins this class at M=1 (2.4-2.6x bf16, ahead of tiled's ~1.1x there).
M > 1 WITH a fused epilogue (the tiled interior store can't apply it), or any M on pre-M5: MATERIALIZE the weight to a transient (N, K) bf16 buffer, then run the dense bf16 MMA (hardware-neutral on pre-M5).

Parameters:

c_type (DType): Output element type (fp16, bf16, fp32). Accumulation is fp32.
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 (AMD's (row, col) contract), threaded to whichever path runs.